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.
Rong-Tian Fu, Dong-Bo Lv, Xue-Ying Wang et al.· Proceedings of the Internati...· 0 citations
Mixture-of-Experts (MoE) has become the de facto architecture for scaling large language models, offering expanded capacity with manageable compute. Pipeline parallelism (PP) is indispensable for distributed MoE training, but state-of-the-art PP schemes face three major limitations: large pipeline bubbles, high per-stage latency due to insufficient overlap of Expert Parallelism (EP) communication, and limited flexibility due to rigid configuration constraints. As a result, training efficiency degrades severely in large-scale clusters. We propose OmniPipe, a flexible bidirectional multi-pipeline parallelism scheme for unified dense and MoE LLM training. It relaxes the rigid constraints of prior bidirectional pipelines by supporting configurable pipeline replica counts and flexible micro-batch scaling. A flexible scheduling strategy further achieves better overlap between EP All-to-All communication and computation, effectively minimizing pipeline bubbles and intra-stage latency to reduce end-to-end pipeline execution time. Implemented within Megatron-LM, OmniPipe enables full 5D hybrid parallelism. On NVIDIA A800 GPU clusters, OmniPipe consistently outperforms the best configurations of existing PP schemes, achieving a geometric-mean 1.12 × speedup on MoE workloads with up to 1.30 × , while also achieving a 1.10 × average speedup on dense models. The results demonstrate that OmniPipe minimizes the pipeline bubble ratio while effectively overlapping EP communication with computation, enabled by the flexible and scalable parallelism scheme of bidirectional pipelines.
Jun Li, Zhi Ma, Shigang Li· Proceedings of the Internati...· 0 citations
We use cookies to run the site and, with your consent, for analytics and to show ads.
See our Cookie Policy.
