Flash-KMeans 在 GPU 上以精确 K-means 比 FAISS 快 200 多倍

一支研究团队发布了 Flash-KMeans，这是一个开源库，在 GPU 上运行精确的 K-means 聚类，速度比 FAISS 快 200 多倍，而且没有改动数学本身。K-means 是最古老、应用最广泛的聚类算法之一，是从向量数据库索引到图像压缩等各种场景背后的主力，而 Flash-KMeans 并不对它做近似，也不改动 Lloyd 的经典流程。相反，它重新思考算法如何在现代 GPU 上搬运数据，正是这种洞见让 FlashAttention 变快：瓶颈不在于算术运算，而在于内存流量。

这一加速来自两个内核级的想法。第一个是 FlashAssign，它把距离计算与在线 argmin 融合在一起，使得分配步骤永远不必把一个巨大的中间距离矩阵写出到内存再读回来，工作始终留在快速的片上内存里。第二个叫 sort-inverse update，它攻克另一个开销很大的步骤，也就是重新计算每个簇的中心。如果做得很朴素，这一步会有成千上万个 GPU 线程争抢着往同样几个累加器里相加，这种高争用模式会拖慢硬件。Flash-KMeans 转而构建一个显式的逆映射，把那些分散的原子写入变成有序、高带宽、分段级的归约。

在一块 NVIDIA H200 上，团队报告端到端速度相比此前最佳基线最高可达 17.9 倍，比 NVIDIA 自家的 cuML 快 33 倍，比 FAISS 快 200 多倍，而 FAISS 正是大多数人在需要大规模做 K-means 时会首先想到的库。而且它是精确的，产生 Lloyd 算法一向会产生的相同聚类，而不是一个悄悄牺牲精度来换取速度的更快近似。这个库用 Triton 编写，也就是那种用于 GPU 内核、类似 Python 的语言，以宽松的 Apache 2.0 协议发布，只需一条 pip 命令即可安装。论文挂在 arXiv，代码放在 GitHub。

这一成果是一个虽小却鲜明的范例，体现了一个在机器学习系统中不断带来回报的模式：最大的收益往往不是来自一个新算法，而是来自重新实现一个旧算法，以尊重它所运行硬件的内存层级。FlashAttention 为 transformer 的注意力做到了这一点，而 Flash-KMeans 为聚类做到了这一点，聚类这个原语位于检索、去重以及支撑当今大量 AI 应用的向量搜索之下。在一个精确、未改动的算法上获得 200 倍加速，并不是边际优化，而是把一个数据集的聚类从几秒钟完成，与等上几分钟之间的区别。诚实的告诫是，那些头条倍数取决于硬件、数据集以及基线是如何配置的，因此现实世界中的收益会有所不同，但方向是清晰的，而且代码就摆在那里供人核验。