Flash-KMeans の紹介：GPU で FAISS よりも 200 倍以上高速に動作する IO 対応型 K-Means アルゴリズム

k-means は数十年にわたりオフラインツールとして使われてきました。データを前処理するために一度実行し、その後次のステップへ進むのです。カリフォルニア大学バークレー校とオースティン大学の研究者チームが Flash-KMeans をリリースしました。これは異なる設定をターゲットにした新しいオープンソースライブラリです。現代の AI パイプラインでは、k-means がトレーニングや推論ループ内で呼び出されるようになりました。その頻度においては、理論上の FLOPs（浮動小数点演算数）よりも、1 回あたりのレイテンシが重要になります。

Flash-KMeans は、標準的な Lloyd の k-means アルゴリズムを IO アウェアに実装したものです。数学的計算を変更するものでもなく、近似を行うものでもありません。GPU 上でアルゴリズムがデータをどのように移動させるかを再構築するだけです。NVIDIA H200 において、研究チームは最良のベースラインと比較して最大で 17.9 倍のエンドツーエンドの高速化を報告しています。NVIDIA cuML（CUDA Machine Learning）との比較では 33 倍、FAISS との比較では 200 倍以上の高速化を報告しています。

Flash-KMeans とは何か

Flash-KMeans は、Triton GPU カーネルで書かれたバッチ処理対応の k-means ライブラリです。Apache 2.0 ライセンスの下で配布され、pip install flash-kmeans でインストールできます。

出力結果は、標準的な Lloyd の k-means と数学的に完全に同一です。高速化の要因は作業を省略することではなく、カーネルレベルでのデータフローにあります。これにより、三角不等式による剪定やコレスセットサンプリングなどのアルゴリズム手法とは明確に区別されます。

標準的な Lloyd 反復には 2 つの段階があります。割り当て段階では、各点からすべてのセントロイドまでの距離を計算し、最も近いものを選択します。更新段階では、各クラスター内の点を平均化して新しいセントロイドを形成します。両方の段階とも単純な算術演算ですが、GPU においてはどちらも計算能力ではなくメモリ帯域幅がボトルネックとなります。

Flash-KMeans が攻撃する 2 つのボトルネック

最初のボトルネックはアサインメント（割り当て）ステージです。標準的なコードでは、高帯域メモリ（HBM）に N×K 形状の距離行列 D を完全に構築します。その後、この行列を書き出し、argmin（最小値インデックスの取得）を実行するために読み戻します。N=65536、K=1024、d=128、B=32 の場合、距離計算自体には 2.6ms かかりますが、行列 D の書き出しと消費には約 23ms を要します。コストの主体は演算ではなく行列そのものです。

Flash-KMeans はこれを FlashAssign で置き換えます。この設計は FlashAttention から着想を得ています。FlashAssign は、ポイント（点）とセントロイド（重心）のタイルを HBM からオンチップ SRAM へストリーミングし、距離計算とオンライン argmin を融合します。完全な N×K 行列が一度に実体化されることはありません。これにより、支配的な IO 複雑度が O(NK) から O(Nd + Kd) に削減されます。カーネルレベルでは FlashAssign は最大で 21.2 倍の性能向上を達成し、あるケースではアサインメント時間を 122.5ms から 5.8ms に短縮しました。

2 つ目のボトルネックはセントロイド更新ステージです。標準的なコードでは散在型（scatter-style）のアトミック加算を使用します。各スレッドがクラスタ ID をキーとして共有和バッファに自身のポイントを追加しますが、多くのスレッドが同時に同じ「ホット」なクラスタにアクセスします。これによりアトミック競合が発生し、ハードウェアによる直列化が生じます。研究チームは H200 において、この部分の有効帯域幅がわずか 50 GB/s であることを測定しました。

Flash-KMeans はこれをソート逆更新（Sort-Inverse Update）に置き換えます。1 次元の割り当てベクトルをクラスター ID で argsort してソートし、同一のクラスター ID が連続するセグメントを形成します。各スレッドブロックはオンチップでセグメントを削減し、その後セグメントごとにアトミック加算（atomic add）を 1 回発行します。重いポイント行列が物理的に並べ替えられることは決してありません。アトミック操作の計算量は O((K + N/B_N)d) に低下します。このカーネルは最大で 6.3 倍の高速化を実現します。

ベンチマーク

研究チームは H200、CUDA 12.8、FP16 データ、次元数 d=128 の環境でテストを行いました。N、K、バッチサイズ B をスweep し、4 つの最適化されたベースラインと比較しました。それらは fast_pytorch_kmeans、fastkmeans、cuML、および FAISS です。

比較報告：速度向上率とワークロードコンテキスト

エンドツーエンド vs 最良ベースライン：最大 17.9 倍（N=8M, K=1024：大規模 N、小規模 K）

vs NVIDIA cuML:33 倍（産業用ライブラリ）

vs FAISS:200 倍以上（産業用ライブラリ）

FlashAssign カーネル：最大 21.2 倍（N=1M, K=8192：割り当て処理）

Sort-Inverse Update カーネル：最大 6.3 倍（N=33M, K=4096：更新処理）

オンメモリ外・大規模スケール：最大 10.5 倍（N=400M, K=16384 vs fastkmeans）

文脈を理解する上で重要な失敗モードが 1 つあります。標準的な PyTorch 実装では、大規模な K のレジームにおいてメモリ不足となり、N×K 行列をマテリアライズ（具体化）することができません。FAISS は、多くの生産環境におけるベクトル検索システムの業界標準ライブラリです。

このライブラリはアウト・オブ・コアでも動作します。10 億点（K=32768, d=128）の場合、1 イテレーションを 41.4 秒で完了し、ベースラインの 261.8 秒と比較して大幅に高速です。計算処理の間に PCIe 転送を隠すために、チャンク化されたストリームのオーバーラップを利用しています。キャッシュを意識したコンパイルヒューリスティックにより、チューニングオーバーヘッドが最大 175 倍削減され、チューニング済み速度に対して誤差は 0.3% 以内となっています。

MTP インタラクティブ・エクスプローラー

#mtp-fk-demo{

background:#111!important;color:#e8e8e8!important;

border:1px solid #222!important;border-radius:14px!important;

font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,Helvetica,Arial,sans-serif!important;

max-width:980px;margin:24px auto!important;padding:0!important;overflow:hidden!important;

line-height:1.5!important

}

#mtp-fk-demo .fk-tab{

background:#181818!important;color:#bdbdbd!important;border:1px solid #262626!important;

border-radius:8px 8px 0 0!important;padding:9px 15px!important;font-size:13px!important;font-weight:600!important;

cursor:pointer!important;transition:all .15s!important

}

#mtp-fk-demo button.fk-b{

background:#76B900!important;color:#0a0a0a!important;border:none!important;border-radius:8px!important;

padding:9px 14px!important;font-size:13px!important;font-weight:700!important;cursor:pointer!important;transition:opacity .15s!important

}

@media(max-width:640px){

}

Marktechpost · Interactive Explainer

Flash-KMeans: GPU メモリを中核に再構築された厳密な k-means アルゴリズム

標準的な k-means と同じ Lloyd の数学（Lloyd's algorithm）を使用 — 高速化の理由はデータフローの最適化のみ。クラスタリングを実行中にライブで確認し、更新時の競合ボトルネックを視覚化し、削減された入出力（IO: Input/Output）の量を算出。

ベストなベースラインと比較してエンドツーエンドで 17.9 倍高速

NVIDIA cuML と比較して 33 倍高速

FAISS と比較して 200 倍以上高速

10 億点のデータポイント、アウトオブコア処理対応

1 · ライブクラスタリング

2 · 更新競合

3 · IO カルキュレーター

データポイント数 (N) 800

クラスタ数 (K) 5

実行

ステップ

新規データ

イテレーション0

重心のシフト—

状態：アイドル

このブラウザ上での 2 次元点に対する実装は、リアルなロイド型 k-means アルゴリズムを実行するものです。このアルゴリズムは Flash-KMeans が加速するものと同じですが、GPU データフローのみが異なります。各ステップは、1 つの割り当てと 1 つの重心更新から構成されます。

プレイボタンを押してください。標準的な散乱・更新方式では、ブロックが同じ「ホット」な重心（赤色）に書き込む際に直列化が発生し、ストールします。一方、ソート逆更新 (Sort-Inverse Update) ではクラスタ ID を先にソートするため、各ブロックは連続するセグメントを 1 つの原子加算でマージでき、競合が発生しません。

タイムライン再生

リセット

標準アトミック演算: O(N·d)

ソート逆アトミック演算 (Sort-Inverse atomics): O((K+N/B)·d)

測定された標準帯域幅：50 GB/s

カーネルの高速化率：6.3 倍

標準的な更新方式では、トークンごとに 1 つの原子加算が発行されます。多くのスレッドが同時に同じ重心にアクセスするため、競合が発生します。クラスタ ID でソートすることで、散乱操作をオンチップメモリ内でのセグメントレベルの還元 (reduction) に変換できます。

標準方式 — N×K 行列を実体化し、計算量は O(NK)—

FlashAssign — 入力をストリーミングし、計算量は O(Nd+Kd)—

—割り当てステップにおける HBM (High Bandwidth Memory) トラフィックの削減（理論値）

ポイント数 N 100 万

クラスタ数 K 1024

次元 d 128

標準的な k-means は、HBM（高帯域幅メモリ）に N×K の距離行列全体を書き込み、その後読み出します。FlashAssign はこれを構築せず、X と C をそれぞれ一度だけ読み取り、割り当てを一度だけ書き込みます。棒グラフは相対的な HBM への往復回数を示しており、FP16（半精度浮動小数点数）での結果です。

© Marktechpost

高速化率：Flash-KMeans 論文 (arXiv:2603.09229)、NVIDIA H200。デモはブラウザ上で動作し、説明用として提供されています · github.com/svg-project/flash-kmeans

/* ---------- パネル 3: I/O 計算機 ---------- */

function fmt(n){

if(n>=1e12)return (n/1e12).toFixed(1)+' TB';

if(n>=1e9)return (n/1e9).toFixed(1)+' GB';

if(n>=1e6)return (n/1e6).toFixed(1)+' MB';

if(n>=1e3)return (n/1e3).toFixed(1)+' KB';

return n.toFixed(0)+' B';

}

function fmtCount(n){

if(n>=1e9)return (n/1e9).toFixed(n>=1e10?0:1)+'B';

if(n>=1e6)return (n/1e6).toFixed(n>=1e7?0:1)+'M';

if(n>=1e3)return (n/1e3).toFixed(n>=1e4?0:1)+'K';

return n+'';

}

function calcIO(){

var N=Math.pow(2,+$('fkIoN').value), K=Math.pow(2,+$('fkIoK').value), d=Math.pow(2,+$('fkIoD').value);

$('fkIoNval').textContent=fmtCount(N);

$('fkIoKval').textContent=fmtCount(K);

$('fkIoDval').textContent=d;

var bpe=2; /* fp16 (半精度浮動小数点) */

var std=2*N*K*bpe; /* 距離行列の書き込みと読み出し */

var fk=(N*d + K*d)*bpe + N*4; /* データ X と中心 C の各 1 回の読み出し + int32 (32 ビット整数) アサインメントの書き込み */

$('fkStdAmt').textContent=fmt(std);

$('fkFkAmt').textContent=fmt(fk);

var ratio=std/fk;

$('fkRatio').textContent=ratio>=10?Math.round(ratio)+'×':ratio.toFixed(1)+'×';

/* バーの幅は、std を 100% とした対数近似スケール上の相対値 */

$('fkStdFill').style.width='100%';

var w=Math.max(3,(fk/std)*100);

$('fkFkFill').style.width=w.toFixed(1)+'%';

}

['fkIoN','fkIoK','fkIoD'].forEach(function(id){$(id).addEventListener('input',calcIO);});

/* 初期化 */

genData();

calcIO();

})();

ユースケース

高速な厳密 k-means は、オフラインだけでなくオンラインでも実行可能な範囲を変えます。

ベクトル検索インデックス：FAISS は k-means を用いて検索インデックスを構築します。k-means が高速化されることで、データが変化した際に一夜明けの再構築ではなく、随時インデックスの再構築が可能になります。

スパースアテンションルーティング：Routing Transformers や Tactic はトークンをクラスタリングしてアテンションをルーティングします。ミリ秒単位の k-means により、推論ループ内での実現が可能になります。

KV キャッシュ圧縮：ClusterKV は意味空間内でトークンをクラスタリングしてキャッシュを圧縮します。低コストなクラスタリングにより、層ごと・ステップごとの圧縮が実用的になります。

低ビット KV 量子化：最近の手法では、KV エントリをコードブックにクラスタリングし、r