Flash-KMeansの紹介：GPU上でFAISSより200倍以上高速なIO認識型の正確なK-Means

Flash-KMeansは、カリフォルニア大学バークレー校とテキサス大学オースティン校の研究チームがリリースしたオープンソースライブラリで、最新のAIパイプラインにおけるk-meansの頻繁な呼び出しに伴うレイテンシボトルネックを解決することを目的としています。従来のk-meansはオフラインツールとして使用されていましたが、現在ではトレーニングや推論ループ内で呼び出されるため、呼び出しごとのレイテンシが理論上のFLOPsよりも重要です。

Flash-KMeansは、標準的なLloydのk-meansのIO認識型実装であり、数学や近似を変更せず、GPU上のデータフローのみを再構築します。NVIDIA H200では、研究チームは最良のベースライン比最大17.9倍、NVIDIA cuML比33倍、FAISS比200倍以上のエンドツーエンド高速化を報告しています。

Flash-KMeansの中心的な革新は、FlashAssignとSort-Inverse Updateの2つの主要カーネルです。FlashAssignはFlashAttentionに着想を得て、点と重心のタイルをHBMからオンチップSRAMにストリーミングし、距離計算とオンラインargminを融合することで、完全なN×K距離行列の実体化を回避します。これにより、割り当てフェーズのIO複雑性がO(NK)からO(Nd+Kd)に削減され、割り当てカーネルは最大21.2倍高速化されます。

Sort-Inverse Updateは、重心更新フェーズのアトミック競合問題に対処します。argsortを使用して割り当てベクトルをクラスタIDでソートし、連続したセグメントを形成します。各スレッドブロックはセグメントをオンチップでリダクションし、セグメントごとに1回のアトミック加算のみを実行します。これにより、更新カーネルは最大6.3倍高速化されます。

ベンチマークはH200上でFP16データ、d=128で実施され、fast_pytorch_kmeans、fastkmeans、cuML、FAISSと比較されました。エンドツーエンドの高速化は最大17.9倍（N=8M、K=1024）、割り当てカーネルは21.2倍、更新カーネルは6.3倍でした。アウトオブコアモードでは、10億ポイント（K=32768、d=128）に対して1反復あたり41.4秒で完了し、ベースラインの261.8秒を大幅に下回りました。

Flash-KMeansには、キャッシュ認識コンパイルヒューリスティックも含まれており、チューニングオーバーヘッドを175倍削減しつつ、最適に近いパフォーマンスを維持します。このライブラリはオープンソースで、Apache 2.0ライセンスの下で提供され、pip install flash-kmeansでインストールできます。

ユースケースとしては、ベクトル検索インデックス構築、スパース注意ルーティング、KVキャッシュ圧縮、低ビットKV量子化、拡散トランスフォーマーなどが挙げられます。APIはfaissやsklearnと類似しており、バッチ処理とマルチGPUをサポートしています。

要約すると、Flash-KMeansはGPUデータフローの最適化により正確なk-meansを大幅に高速化し、頻繁なクラスタリングを必要とする新しいアプリケーションに実用的なツールを提供します。