Flash-GMM: Un Kernel Eficiente en Memoria para Clustering Suave Escalable

Resumen

Presentamos Flash-GMM, un kernel Triton fusionado para el cómputo eficiente de Modelos de Mezcla Gaussianas (GMMs) sobre datos a gran escala en una sola pasada de GPU. Al eliminar la necesidad de materializar la matriz completa de responsabilidades en la memoria de la GPU, Flash-GMM logra una aceleración de 20 veces respecto a implementaciones existentes y permite entrenar conjuntos de datos más de 100 veces mayores de lo que era factible anteriormente en un solo dispositivo. Para demostrar su impacto, integramos Flash-GMM en el cuantizador grueso de IVF para la búsqueda aproximada de vecinos cercanos (ANN). Mostramos que el agrupamiento GMM suave es ahora una sustitución directa viable para k-means, y que las responsabilidades del GMM pueden aprovecharse para asignar vectores frontera a múltiples clústeres. Nuestro enfoque alcanza objetivos fijos de recall con hasta 1.7 veces menos cómputos de distancia, o equivalentemente, produce un incremento de +2 a 12 en recall@10 con un costo computacional equivalente. Publicamos el kernel como proyecto de código abierto.

English

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.