Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering
Abstract
Flash-GMM introduces an efficient fused Triton kernel for Gaussian Mixture Models that achieves significant speedup and enables processing much larger datasets on a single GPU.
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 20times speedup over existing implementations and enables training on datasets more than 100times 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.7times fewer distance computations, or equivalently, yields +2--12 recall@10 at matched computational cost. We release the kernel as an open-source project.
Community
This is an automated message from the Librarian Bot. I found the following papers similar to this paper.
The following papers were recommended by the Semantic Scholar API
- Ascend-RaBitQ: Heterogeneous NPU-CPU Acceleration of Billion-Scale Similarity Search with 1-bit Quantization (2026)
- On Efficient Scaling of GNNs via IO-Aware Layers Implementations (2026)
- DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts (2026)
- Stochastic Sparse Attention for Memory-Bound Inference (2026)
- AdaSplash-2: Faster Differentiable Sparse Attention (2026)
- HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces (2026)
- SAW-INT4: System-Aware 4-Bit KV-Cache Quantization for Real-World LLM Serving (2026)
Please give a thumbs up to this comment if you found it helpful!
If you want recommendations for any Paper on Hugging Face checkout this Space
You can directly ask Librarian Bot for paper recommendations by tagging it in a comment: @librarian-bot recommend
Neat paper. The approach of using a fused Triton kernel to avoid materializing the responsibility matrix seems like a smart way to handle memory bottlenecks. Getting a 20x speedup while scaling up to much larger datasets sounds like a significant jump for GMM training.
I am curious, since you're using this for IVF coarse quantizers, how sensitive is the recall gain to the specific number of clusters when you assign border vectors to multiple cells?
I made a podcast on it with ResearchPod, it makes it easy to get the key concepts on the go:
https://researchpod.app/episode/7614e817-478e-40db-b6da-670fe668c8e5
Get this paper in your agent:
hf papers read 2606.10896 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper