# Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering

arXiv: https://arxiv.org/abs/2606.10896 · HTML: https://mlsanity.com/paper/2606.10896/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 6 of its 27 references that are in the map.
- First posted on arXiv 9 Jun 2026.

## Builds on (most cited first, 6 of 6)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [GMMSeg: Gaussian Mixture based Generative Semantic Segmentation Models](https://mlsanity.com/paper/2210.02025/) (2022)
- [Flash-KMeans: Fast and Memory-Efficient Exact K-Means](https://mlsanity.com/paper/2603.09229/) (2026)
- [RAIRS: Optimizing Redundant Assignment and List Layout for IVF-Based ANN Search](https://mlsanity.com/paper/2601.07183/) (2026)

## Abstract (arXiv)

We present 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 20× speedup over existing implementations and enables training on datasets more than 100× 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× fewer distance computations, or equivalently, yields +2--12 recall@10 at matched computational cost. We release the kernel as an open-source project.