Turbo-GMM speeds up clustering on GPUs with single pass component traversal
Turbo-GMM: FlashAttention-Inspired GPU-Accelerated Gaussian Mixture Model
Machine LearningArtificial IntelligenceHardware ArchitectureDistributed, Parallel, and Cluster Computing
Summary
Clustering data points into groups often uses a method called Gaussian mixture models, which takes a lot of computer time when there are many groups and data points. The authors present Turbo-GMM, a new way to do this clustering faster on GPUs by scanning the groups only once instead of twice. This reduces the time it takes to do the calculations but works best when the number of groups is not too large. In tests, Turbo-GMM was shown to be about twice as fast as a previous GPU method and much faster than common CPU methods.
What this means in practice
- •For data engineers: Accelerate clustering of large datasets with moderate group counts using GPU-optimized Turbo-GMM for faster data processing pipelines.
- •For machine learning platform builders: Integrate Turbo-GMM to improve the performance of clustering services that run on GPU hardware with specialized fused kernels.
Authors
Momin Abbas, Sadia Asif, Asad Ul Haq, Kaoutar El Maghraoui
Abstract
Expectation-maximization (EM) for Gaussian mixture models (GMMs) requires scoring $N$ data points against $K$ mixture components. Prior fused GPU implementations, notably Flash-GMM, eliminate materialization of the $N \times K$ responsibility matrix but traverse the component dimension twice to compute normalization terms and accumulate sufficient statistics. We introduce Turbo-GMM, a single-sweep fused EM kernel that completes responsibility normalization and sufficient-statistic accumulation within one traversal of the component dimension. By retaining the complete component axis on chip, Turbo-GMM avoids the second score computation required by the released two-sweep Flash-GMM implementation. Its applicability is limited by register pressure, occupancy, and the component and feature dimensions, making it particularly suitable for large datasets with relatively small numbers of components. For larger configurations, Turbo-GMM employs a two-pass fallback following the established fused design. We establish algebraic equivalence to standard EM in real arithmetic and evaluate sixteen configurations on an NVIDIA A100. Under the reported benchmarking protocol, Turbo-GMM achieves a median $2.15\times$ speedup over the evaluated Flash-GMM implementation and up to $5.3\times$ and $9.9\times$ over PyTorch and \texttt{scikit-learn}, respectively. The results demonstrate the potential of single-sweep component traversal while identifying on-chip resource constraints as the principal limitation of this approach.