---
title: 'Flash-KMeans: GPU-Optimized k-means'
url: https://www.emergentmind.com/topics/flash-kmeans
type: topic
---

# Flash-KMeans: GPU-Optimized k-means

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 [2603.09229].

## 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 $(i,k)$, then $a_i = \arg\min_{k} D_{ik}$.
- **Update step:** For each cluster $k = 1 \ldots K$,
  \[
  n_k = \sum_{i=1}^N \mathbf{1}\{a_i=k\}, \quad
  s_k = \sum_{i=1}^N \mathbf{1}\{a_i=k\}\,x_i, \quad
  c_k \leftarrow s_k / n_k
  \]

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

- **IO-bound assignment:** Standard GPU implementations materialize the full $N \times K$ distance matrix $D$ in High Bandwidth Memory (HBM), resulting in $O(NK)$ memory writes and reads per iteration. Example: $N=2^{16}, K=2^{10}, d=128$ yields compute time $\sim$2.6 ms but materialization dominates with $\sim$23 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. $>$600 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 $D$. The method tiles both points and centroids, maintaining running $\left(m_i, a_i\right)$ in registers and directly writing final assignments.

- **IO Complexity:** Standard: $O(NK)$ (reads/writes); FlashAssign: $O(Nd + Kd)$ (reads $X$, $C$ once; writes $a$ once).
- **Compute Complexity:** Both unchanged at $O(NKd)$.

This eliminates the dominant $O(NK)$ 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: $O(Nd)$; sort-inverse: $O((K + N/B_N)d)$, where $B_N$ 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$ exceeding GPU RAM limits, data is partitioned into $M$ chunks of $B_N$ points. Data transfer and compute proceed in parallel using two CUDA streams: while one processes chunk $i$, the other asynchronously copies chunk $i+1$. The overall iteration time is
\[
T_{\text{iter}} \approx \max\{T_{\text{copy}}(B_N), T_{\text{assign+update}}(B_N)\}
\]
This effectively hides PCIe transfer overhead as long as $T_{\text{copy}} \leq T_{\text{compute}}$.

### 3.2 Cache-Aware Compile Heuristic

To avoid runtime auto-tuning overheads, tile sizes $(B_N, B_K)$ are selected via direct calculation from L1/L2 on-chip buffer capacities:
\[
B_N d \times 4 \leq M_{L1}/2, \quad B_K d \times 4 \leq M_{L2}/8
\]
Empirical tuning shows $<0.3\%$ suboptimality relative to exhaustive search, with compile/search times reduced by up to $175\times$.

## 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 $(N,K,d)$         | Best Baseline (ms) | cuML (ms) | FAISS (ms) | Flash-KMeans (ms) | Speedup vs Best Baseline |
|----------------------------|--------------------|-----------|------------|-------------------|-------------------------|
| $(8\times10^6, 1024, 128)$ | 170.2              | 192.9     | 8517       | **9.5**           | 17.9×                   |
| $(10^6, 64\mathrm{K}, 512)$| 118.4              | 137.5     | 27,500     | **21.8**          | 5.4×                    |
| $(4\times10^8, 16,384, 128)$† | 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 $K$ workloads.

## 5. Hardware Requirements and Trade-Offs

- **Hardware:** Requires HBM2(e) GPU with $\geq80$ GB and compute capability $\geq9.x$ (to enable async prefetch).
- **Memory footprint:** Supports out-of-core data up to $N=10^9$ using chunked streaming; peak VRAM load is $2 \times B_N \cdot d$ floats.
- **Positive trade-offs:** Provides mathematically exact $k$-means (no approximation), eliminates major IO and atomic-contention sources.
- **Costs:** Introduces an additional argsort operation $(O(N\log N))$ per iteration (highly optimized on GPU, typically small); requires use of custom kernels for full integration.

In sum, Flash-KMeans aligns the $k$-means algorithmic workflow with the realities of modern GPU architectures, reducing IO from $O(NK)$ to $O(Nd+Kd)$, minimizing atomic operations, and achieving significant empirical speedups for both in-core and out-of-core workloads, all while retaining algorithmic exactness [2603.09229].

Source: https://www.emergentmind.com/topics/flash-kmeans