K-means is often treated as an offline preprocessing primitive: cluster embeddings, build codebooks, initialize vector partitions, or group examples before another system consumes the result. Modern AI systems increasingly want clustering inside online or near-online paths, such as dynamic routing, streaming dataset organization, vector-index maintenance, and semantic batching. At that scale, the bottleneck is not the mathematical definition of Lloyd’s algorithm; it is how the assignment and update steps move data through GPU memory.
Flash-KMeans reframes exact k-means as an IO-aware GPU kernel problem. Traditional GPU implementations can materialize an N x K distance matrix in high-bandwidth memory, then read it back to find the nearest centroid for each point. They can also update centroids with many irregular atomic writes. Flash-KMeans avoids both costs: FlashAssign fuses distance computation with online argmin so the full distance matrix is never written, and the sort-inverse update transforms scattered atomic accumulation into localized reductions.
The broader pattern is useful beyond this specific project. Many classical algorithms become practical online primitives when intermediate tensors are treated as a liability, reductions are fused into producers, and irregular writes are rearranged into memory-friendly segments.
The assignment stage computes the nearest centroid for every point. A naive GPU formulation computes distances for all point-centroid pairs and stores them before selecting the minimum. That storage can dominate runtime and memory footprint because the intermediate matrix grows with both sample count and cluster count.
An online argmin keeps only the best distance and best centroid seen so far for each point. Distances are computed tile by tile, compared immediately, and discarded. This turns the assignment stage from a “write all distances, then reduce” pipeline into a fused “compute and keep the winner” kernel.
The update stage sums all points assigned to each centroid. If every point atomically adds into its centroid’s accumulator, popular clusters create write contention. A sort-inverse update groups points by cluster and performs segment-level reductions, trading random atomics for contiguous memory work.
The k-means objective is unchanged. The speedup comes from changing where intermediate values live, when reductions happen, and whether writes are scattered or grouped.
The assignment kernel streams centroid tiles through fast on-chip memory, computes partial distances, and maintains the current nearest centroid. Only final labels and minimal distances need to be written back. This mirrors the IO-aware design family popularized by FlashAttention: exact results, but with tiled memory movement and producer-side reduction.
for point_tile in points:
best_dist = [+inf for point in point_tile]
best_id = [-1 for point in point_tile]
for centroid_tile in centroids:
distances = pairwise_l2(point_tile, centroid_tile)
best_dist, best_id = online_argmin(
best_dist,
best_id,
distances,
centroid_tile.ids,
)
write(assignments=best_id, min_distances=best_dist)After assignment, the runtime constructs a mapping from centroid IDs back to the points assigned to them. Once points are grouped, each centroid update becomes a local reduction over a segment. This avoids the worst case where thousands of threads contend on the same centroid accumulator.
Large inputs may not fit in GPU memory. The implementation handles CPU-resident tensors by copying chunks from pinned host memory to the GPU. A double-buffered stream can prefetch the next chunk while the current chunk is computing, keeping the GPU busy when PCIe bandwidth allows.
km = FlashKMeans(d=128, k=8192, niter=100)
labels = km.fit_predict(large_cpu_tensor)
# Internally:
# - split N points into chunks
# - copy chunk i to GPU
# - compute local assignment and partial sums
# - prefetch chunk i+1 while reducing chunk i
# - combine partial centroid sums and countsLarge embedding corpora are often clustered before building an ANN index, selecting exemplars, deduplicating records, or planning balanced shards. Faster exact clustering reduces pressure to downsample or switch to approximate clustering before measuring quality.
Partitioned vector indexes depend on centroid assignment. If clustering can be recomputed or refreshed quickly, systems can adapt partitions as data distributions shift rather than treating clustering as a rare offline rebuild.
Routing transformers, sparse video generation, mixture systems, and semantic batching can use clustering-like operations to group tokens or examples. IO-aware exact kernels make those routing decisions more plausible inside performance-sensitive loops.
Choose tile sizes that fit the GPU’s memory hierarchy and expose enough parallelism. The point tile controls batch work per block; the centroid tile controls reuse. The product should avoid spilling large intermediate distance blocks to global memory.
A production library should report when Triton or CUDA kernels are unavailable and fall back to a chunked PyTorch implementation only when acceptable. Silent fallback can hide performance regressions and turn a memory-safe workload into an out-of-memory failure.
if use_triton and cuda_available():
labels, centroids = triton_kmeans(x, k, max_iters=niter)
else:
labels, centroids = torch_chunked_kmeans(
x,
k,
chunk_size_N=32768,
chunk_size_K=1024,
)For very large N, split contiguous point ranges across GPUs. Each GPU runs assignment and local accumulation over its partition. Then gather small centroid sums and counts, reduce them, broadcast updated centroids, and overlap the next host-to-device transfer with the reduce path when hardware links allow.
N x K matrix can be much larger than inputs, centroids, labels, and final distances. Eliminating that write-read cycle is the primary win.K, dtype, and batch shape. Keep tuning parameters observable.