The flash-kmeans project presents an IO-aware, batched K-Means clustering algorithm implemented with highly optimized Triton GPU kernels. This library is designed to offer significant performance improvements for large-scale clustering tasks, particularly those involving high-dimensional or massive datasets, and serves as the official K-Means implementation for Sparse VideoGen2.

The core methodology revolves around leveraging NVIDIA Triton, a programming language for writing highly efficient GPU kernels, to accelerate the computationally intensive parts of K-Means clustering, primarily the assignment step. The implementation is "IO-aware," meaning it efficiently manages data transfer between CPU and GPU memory and within GPU memory hierarchies to minimize bottlenecks, often employing double-buffered streaming.

Key technical aspects and optimizations include:

1. Batched Processing: The algorithm processes data in batches, allowing for efficient utilization of GPU resources and reducing memory overhead compared to non-batched approaches. The primary function batch_kmeans_Euclid takes input data x of shape $(B, N, D)$, where $B$ is batch size, $N$ is the number of data points per batch, and $D$ is the dimensionality.

2. Triton Kernel Design for Assignment (E-step): The assignment kernel, responsible for finding the closest centroid for each data point, is meticulously optimized and dispatches between two main flavors based on input shape and data type:
   • Small-D Path ($D \leq 512$): For relatively lower-dimensional data, this kernel loads a tile of data points $x_{tile}$ of shape $(BN, D)$ once. It then streams over the centroids, processing them in chunks of size BLOCK_K. Distances are accumulated without materializing the full $(B, N, K)$ distance matrix, preserving a "K-streaming" property that minimizes memory footprint. Performance is tuned using per-architecture heuristics (e.g., for H200, H100, A100, GB10 GPUs) to select optimal tile sizes (BLOCK_N, BLOCK_K), num_warps, and num_stages.
   • Split-D Path ($D > 512$ or Shared Memory Constraints): For higher-dimensional data or when the small-D kernel exceeds shared memory limits, this path employs an outer loop over centroids (K) and an inner loop over dimensions (D), tiling the D-loop by BLOCK_D. Crucially, the cross-accumulator (which holds partial sums for distance calculations for a tile of $BN$ data points against a tile of $BK$ centroids) is maintained in fast GPU registers across the D-loop. This design also preserves the K-streaming property, avoiding the full distance matrix. For unknown GPU architectures, a conservative fallback configuration ($BLOCK_N=32, BLOCK_K=32, BLOCK_D=32, num_warps=4, num_stages=1$) is used, dispatching to this path to guarantee shared memory fit, though performance may be suboptimal without specific tuning.

3. Large Tensor Handling: For datasets that exceed GPU memory capacity, flash-kmeans supports large $N$ by transferring data from CPU pinned memory to the GPU in chunks. Both flash-kmeans and compared baselines (e.g., fastkmeans) adopt this strategy, ensuring scalability for massive datasets.

4. Multi-GPU Support (kmeans_largeN): The library offers automatic multi-GPU scaling for large $N$ workloads.
   • Data Partitioning: The total $N$ data points are partitioned into contiguous blocks, with one partition assigned to each available GPU.
   • Independent Pipelines: Each GPU independently executes its own double-buffered Host-to-Device (H2D) data transfer and computation pipeline over its assigned data partition. This linear scaling of PCIe bandwidth with GPU count is achieved because data transfers occur in parallel.
   • Lightweight AllReduce: After each GPU completes its local accumulation of partial centroid sums (approximately 4 MB) and counts (approximately 32 KB), these partial results are gathered to GPU 0 via Device-to-Device (D2D) copies over NVLink. GPU 0 then performs the reduction, and the updated centroids are broadcast back to all participating GPUs. This manual gather-reduce-broadcast mechanism is utilized instead of NCCL, as the small data size makes manual copies faster and avoids external library dependencies, keeping the process contained within a single application.
   • H2D / D2D Overlap: To hide communication latency, the H2D transfers (which use PCIe) and the AllReduce operation (which uses NVLink) are overlapped. The first H2D block for the *next* iteration is prefetched concurrently with the AllReduce operation of the *current* iteration, effectively masking the latency of centroid synchronization.

This comprehensive approach allows flash-kmeans to achieve significant performance improvements over existing K-Means implementations, particularly on NVIDIA GPUs and for large-scale, high-dimensional data.