Flash-KMeans 讓精確 K-means 在 GPU 上比 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 撰寫，這是一種類 Python、用於 GPU 核心的語言，並依寬鬆的 Apache 2.0 授權發布，只需單一一行 pip 指令即可安裝。論文公開在 arXiv 上，程式碼則在 GitHub 上。

這個成果是一個小而鮮明的範例，呼應了一個在機器學習系統中不斷帶來回報的模式：最大的勝利往往不是來自新演算法，而是來自重新實作一個舊演算法，以尊重其運行硬體的記憶體階層。FlashAttention 為 transformer 的注意力機制做到了這點，而 Flash-KMeans 為分群做到了這點，這個基本元件支撐著檢索、去重，以及驅動當今眾多 AI 應用的向量搜尋。在一個精確、未更動的演算法上達成 200 倍加速，並非邊際優化，而是把一個資料集在數秒內分群完成、與枯等數分鐘之間的差別。誠實的但書是，這些亮眼的倍數取決於硬體、資料集，以及基準如何設定，因此實際收益會有所不同，但方向很清楚，而且程式碼就擺在那裡可供查證。