# Communication-Avoiding Linear Algebraic Kernel K-Means on GPUs

arXiv: https://arxiv.org/abs/2601.17136 · HTML: https://mlsanity.com/paper/2601.17136/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 62 references that are in the map.
- First posted on arXiv 23 Jan 2026.

## Builds on (most cited first, 19 of 19)

- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [Reducing Communication in Graph Neural Network Training](https://mlsanity.com/paper/2005.03300/) (2020)
- [Randomly pivoted Cholesky: Practical approximation of a kernel matrix with few entry evaluations](https://mlsanity.com/paper/2207.06503/) (2022)
- [Distributed k-Means and k-Median Clustering on General Topologies](https://mlsanity.com/paper/1306.0604/) (2013)
- [Red-blue pebbling revisited: near optimal parallel matrix-matrix multiplication](https://mlsanity.com/paper/1908.09606/) (2019)
- [Scalable Kernel K-Means Clustering with Nystrom Approximation: Relative-Error Bounds](https://mlsanity.com/paper/1706.02803/) (2017)
- [Scaling betweenness centrality using communication-efficient sparse matrix multiplication](https://mlsanity.com/paper/1609.07008/) (2016)
- [Communication-Efficient Jaccard Similarity for High-Performance Distributed Genome Comparisons](https://mlsanity.com/paper/1911.04200/) (2019)
- [Optimizing CUDA Code By Kernel Fusion---Application on BLAS](https://mlsanity.com/paper/1305.1183/) (2013)
- [Distributed Matrix-Based Sampling for Graph Neural Network Training](https://mlsanity.com/paper/2311.02909/) (2023)
- [A Study of Clustering Techniques and Hierarchical Matrix Formats for Kernel Ridge Regression](https://mlsanity.com/paper/1803.10274/) (2018)
- [Parallel String Graph Construction and Transitive Reduction for De Novo Genome Assembly](https://mlsanity.com/paper/2010.10055/) (2020)
- [Combinatorial BLAS 2.0: Scaling combinatorial algorithms on distributed-memory systems](https://mlsanity.com/paper/2106.14402/) (2021)
- [Extreme-scale many-against-many protein similarity search](https://mlsanity.com/paper/2303.01845/) (2023)
- [RDMA-Based Algorithms for Sparse Matrix Multiplication on GPUs](https://mlsanity.com/paper/2311.18141/) (2023)
- [A sparsity-aware distributed-memory algorithm for sparse-sparse matrix multiplication](https://mlsanity.com/paper/2408.14558/) (2024)
- [Distributed Kernel K-Means for Large Scale Clustering](https://mlsanity.com/paper/1710.03013/) (2017)
- [Distributed-Memory Parallel Contig Generation for De Novo Long-Read Genome Assembly](https://mlsanity.com/paper/2207.04350/) (2022)
- [Popcorn: Accelerating Kernel K-means on GPUs through Sparse Linear Algebra](https://mlsanity.com/paper/2501.05587/) (2025)

## Abstract (arXiv)

Clustering is an important tool in data analysis, with K-means being popular for its simplicity and versatility. However, it cannot handle non-linearly separable clusters. Kernel K-means addresses this limitation but requires a large kernel matrix, making it computationally and memory intensive. Prior work has accelerated Kernel K-means by formulating it using sparse linear algebra primitives and implementing it on a single GPU. However, that approach cannot run on datasets with more than approximately 80,000 samples due to limited GPU memory. In this work, we address this issue by presenting a suite of distributed-memory parallel algorithms for large-scale Kernel K-means clustering on multi-GPU systems. Our approach maps the most computationally expensive components of Kernel K-means onto communication-efficient distributed linear algebra primitives uniquely tailored for Kernel K-means, enabling highly scalable implementations that efficiently cluster million-scale datasets. Central to our work is the design of partitioning schemes that enable communication-efficient composition of the linear algebra primitives that appear in Kernel K-means. Our 1.5D algorithm consistently achieves the highest performance, enabling Kernel K-means to scale to data one to two orders of magnitude larger than previously practical. On 256 GPUs, it achieves a geometric mean weak scaling efficiency of 79.7% and a geometric mean strong scaling speedup of 4.2×. Compared to our 1D algorithm, the 1.5D approach achieves up to a 3.6× speedup on 256 GPUs and reduces clustering time from over an hour to under two seconds relative to a single-GPU sliding window implementation. Our results show that distributed algorithms designed with application-specific linear algebraic formulations can achieve substantial performance improvement.