svg-project/flash-kmeans

What it solves

Flash-KMeans provides a high-performance, memory-efficient implementation of the K-Means clustering algorithm. It addresses the common problem of Out-of-Memory (OOM) errors and slow computation speeds when dealing with massive datasets (large N) or high-dimensional data (large D) on GPUs, avoiding the materialization of large distance matrices.

How it works

The project uses Triton GPU kernels to implement an IO-aware batched K-Means. It employs two primary execution paths based on the data dimensions:

• Small-D path: Optimized for dimensions $\le 512$, using hand-tuned heuristics for specific GPU architectures (H200, H100, A100, GB10).
• Split-D path: Used for dimensions $> 512$ or when shared memory is limited, tiling the dimension loop to preserve the K-streaming property.

For datasets too large for a single GPU, it implements a double-buffered streaming design that transfers data from CPU to GPU in chunks. It also supports multi-GPU scaling by partitioning data across GPUs and using a lightweight manual gather-reduce-broadcast mechanism for centroid updates, avoiding NCCL dependencies.

Who it’s for

It is designed for researchers and developers working with large-scale clustering tasks, specifically those implementing systems like Sparse VideoGen2, or anyone needing a fast, exact K-Means implementation that scales across multiple GPUs.

Highlights

• Triton-based acceleration: Significant performance gains over standard PyTorch and other Triton implementations.
• Memory efficiency: Prevents OOM by avoiding the creation of the full distance matrix.
• Automatic dispatch: Automatically switches between Small-D and Split-D kernels based on input shape and dtype.
• Multi-GPU scaling: Linear PCIe bandwidth scaling and overlap of H2D transfers with centroid reduction.
• Broad hardware support: Includes tuned configs for modern NVIDIA GPUs with a conservative fallback for unknown architectures.