Search papers, labs, and topics across Lattice.
Flash-GMM introduces a novel fused Triton kernel that significantly enhances the computation of Gaussian Mixture Models (GMMs) by avoiding the full responsibility matrix in GPU memory, resulting in a 20脳 speedup. This advancement allows for the training of GMMs on datasets over 100脳 larger than previously possible on a single GPU, making soft clustering more scalable. When integrated into the IVF coarse quantizer for approximate nearest-neighbor search, Flash-GMM demonstrates that it can effectively replace k-means clustering while achieving improved recall with fewer distance computations.
Achieving a 20脳 speedup in GMM computation opens the door for soft clustering on datasets 100脳 larger than before, revolutionizing scalable machine learning.
We present \textbf{Flash-GMM}, a fused Triton kernel for efficient computation of Gaussian Mixture Models (GMMs) over large-scale data in a single GPU pass. By eliminating the need to materialize the full responsibility matrix in GPU memory, Flash-GMM achieves a \textbf{20$\times$} speedup over existing implementations and enables training on datasets more than \textbf{100$\times$} larger than previously feasible on one device. To demonstrate its impact, we integrate Flash-GMM into the IVF coarse quantizer for approximate nearest-neighbor (ANN) search. We show that soft GMM clustering is now a viable drop-in replacement for $k$-means, and that GMM responsibilities can be leveraged to assign border vectors to multiple clusters. Our approach reaches fixed recall targets with up to $1.7\times$ fewer distance computations, or equivalently, yields $+2$--$12$ recall@10 at matched computational cost. We release the kernel as an open-source project.
