Flash-KMeans: GPU-Optimized k-means

• Flash-KMeans is a GPU-native, IO-aware k-means clustering method that transforms the algorithm into an online component by eliminating IO bottlenecks and atomic contention.
• It introduces kernel-level innovations such as FlashAssign to fuse distance computation with argmin selection and sort-inverse update to replace atomic contention with efficient segment-level reductions.
• The approach integrates algorithm-system co-designs like chunked-stream overlap and cache-aware heuristics to boost out-of-core scalability and deliver significant empirical speedups over existing methods.

Flash-KMeans is a GPU-native, IO-aware, and contention-free implementation of $k$-means clustering, designed to transform the algorithm from an offline preprocessing primitive into a deployable online component. By reorganizing both kernel logic and system-level dataflow, Flash-KMeans overcomes persistent bottlenecks in GPU-based $k$-means—specifically, the IO amplification arising from distance matrix materialization and the bandwidth collapse due to atomic write contention in centroid updates. Flash-KMeans introduces two core kernel-level innovations—FlashAssign and sort-inverse update—alongside algorithm-system co-designs such as chunked-stream overlap and cache-aware compile heuristics, enabling both substantial speedups and robust out-of-core scalability on modern GPU hardware (Yang et al., 10 Mar 2026).

1. Problem Formulation and GPU Bottlenecks

Let $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$ denote the dataset of $N$ points in $d$ dimensions, and $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$ the $K$ cluster centroids. The classical Lloyd’s formulation of the $k$-means objective is: $$\min_{C,\, a \in \{1,\dots,K\}^N} \sum_{i=1}^N \| x_i - c_{a_i} \|_2^2$$ Each Lloyd iteration alternates between:

• Assignment step: Compute $D_{ik} = \|x_i - c_k\|_2^2$ for all $k$0, then $k$1.
• Update step: For each cluster $k$2,

$k$3

On current GPUs (e.g., NVIDIA H200), two primary bottlenecks emerge:

• IO-bound assignment: Standard GPU implementations materialize the full $k$4 distance matrix $k$5 in High Bandwidth Memory (HBM), resulting in $k$6 memory writes and reads per iteration. Example: $k$7 yields compute time $k$82.6 ms but materialization dominates with $k$923 ms.
• Atomic-contention update: Centroid updates depend on scatter-style atomic writes, leading to serialization at "hot" centroids and write bandwidth collapse (e.g., 50 GB/s observed vs. $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$0600 GB/s achievable in regular reductions).

2. Kernel-Level Innovations

2.1 FlashAssign: Fused Distance and Argmin

FlashAssign fuses pairwise distance computation with online argmin selection, entirely bypassing the need to explicitly materialize $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$1. The method tiles both points and centroids, maintaining running $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$2 in registers and directly writing final assignments.

• IO Complexity: Standard: $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$3 (reads/writes); FlashAssign: $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$4 (reads $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$5, $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$6 once; writes $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$7 once).
• Compute Complexity: Both unchanged at $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$8.

This eliminates the dominant $$X = [x_1,\dots,x_N]^\top \in \mathbb{R}^{N \times d}$$9 IO penalty in the assignment phase.

2.2 Sort-Inverse Update: Segment-Level Reductions

The sort-inverse update replaces per-token scatter-atomic operations with reductions over cluster-wise segments. After assignments, it applies an argsort over assignment keys, forming contiguous runs per cluster, which are then locally reduced and aggregated with minimal atomic contention.

• Atomic Op Count: Standard: $N$0; sort-inverse: $N$1, where $N$2 is the tile/chunk size.
• Bandwidth Impact: Restores full reduction bandwidth by transforming irregular atomics to segment-level localized reductions, eliminating hot-spot centroids as bottlenecks.

3. Algorithm-System Co-Design

To ensure practical deployability for massive datasets and variable hardware constraints, Flash-KMeans integrates system-level designs:

3.1 Chunked-Stream Overlap

For $N$3 exceeding GPU RAM limits, data is partitioned into $N$4 chunks of $N$5 points. Data transfer and compute proceed in parallel using two CUDA streams: while one processes chunk $N$6, the other asynchronously copies chunk $N$7. The overall iteration time is

$N$8

This effectively hides PCIe transfer overhead as long as $N$9.

3.2 Cache-Aware Compile Heuristic

To avoid runtime auto-tuning overheads, tile sizes $d$0 are selected via direct calculation from L1/L2 on-chip buffer capacities: $d$1 Empirical tuning shows $d$2 suboptimality relative to exhaustive search, with compile/search times reduced by up to $d$3.

4. Empirical Performance

Flash-KMeans is benchmarked against fast_pytorch_kmeans, fastkmeans, NVIDIA cuML, and FAISS. All tests are run as single Lloyd iteration latencies on H200 GPUs with float32 data. Key results:

Workload $d$4	Best Baseline (ms)	cuML (ms)	FAISS (ms)	Flash-KMeans (ms)	Speedup vs Best Baseline
$d$5	170.2	192.9	8517	9.5	17.9×
$d$6	118.4	137.5	27,500	21.8	5.4×
$d$7†	88,400 s	—	—	8.4 s	10.5×

† Out-of-core comparison vs. fastkmeans only.

Additional kernel-level speedups:

• FlashAssign: up to 21.2×
• Sort-Inverse Update: up to 6.3×

Key summary: Flash-KMeans achieves up to 17.9× speedup vs. fast_pytorch_kmeans, 33× vs. cuML, and over 200× vs. FAISS on large $d$8 workloads.

5. Hardware Requirements and Trade-Offs

• Hardware: Requires HBM2(e) GPU with $d$9 GB and compute capability $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$0 (to enable async prefetch).
• Memory footprint: Supports out-of-core data up to $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$1 using chunked streaming; peak VRAM load is $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$2 floats.
• Positive trade-offs: Provides mathematically exact $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$3-means (no approximation), eliminates major IO and atomic-contention sources.
• Costs: Introduces an additional argsort operation $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$4 per iteration (highly optimized on GPU, typically small); requires use of custom kernels for full integration.

In sum, Flash-KMeans aligns the $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$5-means algorithmic workflow with the realities of modern GPU architectures, reducing IO from $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$6 to $$C = [c_1,\dots,c_K]^\top \in \mathbb{R}^{K \times d}$$7, minimizing atomic operations, and achieving significant empirical speedups for both in-core and out-of-core workloads, all while retaining algorithmic exactness (Yang et al., 10 Mar 2026).