Flash-GMM: Um Kernel Eficiente em Memória para Agrupamento Suave Escalável

Resumo

Apresentamos o Flash-GMM, um kernel Triton fundido para o cálculo eficiente de Modelos de Misturas Gaussianas (GMMs) em dados de grande escala em uma única passagem pela GPU. Ao eliminar a necessidade de materializar a matriz completa de responsabilidades na memória da GPU, o Flash-GMM atinge uma aceleração de 20 vezes em relação às implementações existentes e permite o treinamento em conjuntos de dados mais de 100 vezes maiores do que anteriormente viável em um único dispositivo. Para demonstrar seu impacto, integramos o Flash-GMM ao quantizador grosseiro IVF para busca aproximada de vizinhos mais próximos (ANN). Mostramos que o agrupamento suave por GMM é agora uma substituição direta viável para o k-means, e que as responsabilidades do GMM podem ser aproveitadas para atribuir vetores de fronteira a múltiplos clusters. Nossa abordagem atinge alvos fixos de recall com até 1,7 vezes menos cálculos de distância, ou, equivalentemente, produz +2–12 recall@10 com custo computacional equivalente. Disponibilizamos o kernel como um projeto de código aberto.

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.