Meet Flash-KMeans: An IO-Aware, Exact K-Means That Runs Over 200× Faster Than FAISS on GPUs

📌 【UC Berkeley & UT Austin 最新研究】Flash-KMeans：讓 GPU 上的 K-Means 速度提升 200 倍以上

K-Means 演算法已經存在數十年，過去我們習慣將其視為一種「離線」的預處理工具：跑一次，處理完數據，然後就結束。但隨著現代 AI 管道的演進，K-Means 越來越常被直接嵌入到訓練（Training）與推論（Inference）的迴圈中。

當 K-Means 變成頻繁調用的組件時，決定效能的不再是理論上的運算量 (FLOPs)，而是單次調用的延遲 (Latency)。

🤔 運算很快，但數據搬運太慢

在 GPU 上執行標準的 Lloyd’s K-Means 時，瓶頸其實不在於數學運算，而是在於記憶體 I/O。

以典型的「指派階段 (Assignment stage)」為例：程式需要計算每個點與所有質心的距離，並找出最近的一個。傳統做法會先在高頻寬記憶體 (HBM) 中建立一個巨大的距離矩陣 $D$（維度為 $N \times K$），寫入矩陣後再讀回來執行 $\text{argmin}$。

研究數據顯示，在特定規模下（$N=65536, K=1024, d=128, B=32$），實際的數學運算僅需 2.6ms，但寫入與讀取矩陣 $D$ 卻耗時 23ms。這意味著大部分的時間都浪費在「搬運數據」而非「計算」。

🧪 借鑒 FlashAttention 的 IO-Aware 設計

為了打破這個瓶頸，來自 UC Berkeley 與 UT Austin 的研究團隊開發了 Flash-KMeans。其核心創新在於 FlashAssign，其設計理念與著名的 FlashAttention 相似：

• 減少 HBM 讀寫：不再建立巨大的中間距離矩陣。
• SRAM 串流處理：將點與質心的分塊 (Tiles) 從 HBM 串流到晶片內的高速 SRAM 中。
• 算子融合 (Operator Fusion)：將距離計算與 $\text{argmin}$ 過程直接融合在同一個 Kernel 中同步完成。

最關鍵的是，Flash-KMeans 沒有改變數學邏輯，也沒有使用近似算法。它不採取像三角不等式剪枝 (Triangle-inequality pruning) 或核心集採樣 (Coreset sampling) 等減少工作量的做法，其輸出結果與標準 Lloyd’s K-Means 完全一致。

🚀 對比 FAISS 速度提升超過 200 倍

在 NVIDIA H200 GPU 的測試中，Flash-KMeans 展現了極其顯著的端到端加速效果：

• 相較於 FAISS：速度提升超過 200 倍
• 相較於 NVIDIA cuML：速度提升 33 倍
• 相較於最佳基準線 (Best baseline)：最高提升 17.9 倍

這種量級的提升，讓 K-Means 真正能以極低延遲地整合進高效能的 AI 訓練流程中。

💡 從「減少計算」轉向「優化數據流」

這次的研究再次證明了一個重要的工程洞察：在現代 GPU 硬體上，許多演算法的瓶頸在於記憶體帶寬而非運算能力。透過重新設計數據在 HBM 與 SRAM 之間的流動路徑，即使不減少運算量，也能獲得數十倍甚至數百倍的效能增益。

🎯 實務啟示：AI 管道優化的新方向

如果你正在構建需要頻繁執行分群 (Clustering) 的大規模訓練管道，Flash-KMeans 提供了一個高效且精準的選擇。由於它使用 Triton 撰寫且採取 Apache 2.0 開源，開發者可以快速整合進現有流程。

• 安裝簡單：pip install flash-kmeans
• 精準度保證：無需在「速度」與「精確度」之間做 trade-off，因為它是 Exact K-Means。

🔗 相關資訊 📝 Meet Flash-KMeans: An IO-Aware, Exact K-Means That Runs Over 200× Faster Than FAISS on GPUs 👤 Asif Razzaq 🔗 詳情：https://www.marktechpost.com/2026/06/15/meet-flash-kmeans-an-io-aware-exact-k-means-that-runs-over-200x-faster-than-faiss-on-gpus/

你的 AI 管道中是否有記憶體 I/O 導致的瓶頸？歡迎在下方討論你的優化經驗 👇

#AI #GPU #KMeans #DeepLearning #Triton #CUDA #PerformanceOptimization #UCBerkeley #UTAustin