MX-KMeans: Accelerating K-Means Clustering through Microscaling Quantization

Abstract

K-Means clustering is a classical unsupervised learning method widely used for its simplicity, efficiency, and broad applicability. In this work, we first analyze the numerical distributions of representative K-Means datasets and identify an opportunity for low-precision acceleration through hardware-native microscaling formats. Based on this observation, we propose MX-KMeans, a high-performance K-Means solution that accelerates clustering with precision-adaptive microscaling quantization. In practice, our MX-KMeans assigns different low-precision formats to data points in a pointwise manner according to their tolerance for quantization error: points well inside a cluster can use extremely low precision, while boundary points are protected with higher precision. The dominant point-to-centroid distance computation is then executed as mixed-precision GEMM, where NVIDIA Blackwell Tensor Cores directly consume microscaled low-precision values and their scaling factors without explicit dequantization overhead. Furthermore, MX-KMeans can be extended beyond standard Lloyd-style K-Means by reformulating pruning-based variants, including Elkan, Hamerly, and Yinyang, as a unified masked distance computation and further lowering it to Tensor-Core-friendly dense GEMM operations. Extensive experiments on an NVIDIA RTX 5090 GPU show that MX-KMeans achieves up to 4.08 × and 75.82 × end-to-end speedups over industry-standard cuML and FAISS, respectively, while preserving the clustering quality of high-precision baselines.