Comment by matrix2596

looks like flash attention concepts applied to kmeans, nice speedup results