looks like flash attention concepts applied to kmeans, nice speedup results
wood_spirit•1h ago
Does this have corresponding speed ups or memory gains for normal CPUs too? Just thinking about all the cups of coffee that have been made and drunk while scikit-learn kmeans chugs through a notebook :)
snovv_crash•41m ago
For CPU with bigger K you would put the centroids in a search tree, so take advantage of the sparsity, while a GPU would calculate the full NxK distance matrix. So from my understanding the bottleneck they are fixing doesn't show up on CPU.
xavxav•19m ago
search trees tend not to scale well to higher dimensions though, right?
from what I've seen I had the impression that Yinyang k-means was the best way to take advantage of the sparsity.
matrix2596•1h ago