llama.cpp b12191、MoEキャッシュミスの一部をCPUで計算

llama.cppはリリースb12191(#30966)で、MoEキャッシュミスのうち指定した割合をGPUへアップロードせずCPU上のホストエキスパートで計算する--moe-cache-cpu-share Fと、キャッシュ引数を選ぶmoe-cache-tune.pyを追加しました。llama.cppはggml-orgが公開するC/C++製の推論ソフトウェアで、CPUと各社GPUの併用に対応します。

何が新しいか

これまでMoEキャッシュのミスはデバイスへのアップロードで処理していましたが、今回の変更でミスの一部をCPUで計算できるようになりました。グラフはエキスパートFFNを2回構築します。

• デバイス側: キャッシュスロットから計算し、CPUに回したミスはゼロスロットを読む
• CPU側: ホストエキスパートから計算し、デバイスで計算したエキスパートはid -1になる
• 2つの結果はダウンバイアスの前に加算する

ggml-cpuのMUL_MAT_IDは負のidをスキップしてゼロ行を書き出します。CPUミスの選択はランダムではなくCRFで行います(#31853)。

動作条件と内部処理

項目	内容
CPUシェア開始	キャッシュグループ内の全スロットがエキスパートを保持した後
空キャッシュ時	キャッシュが埋まるまでミスをすべてアップロード
CPU部分の配置	ホストエキスパートを読むため、スケジューラがCPUに配置
ゼロスロット	各トークンが異なるゼロスロットを使用
制約	1トークンが同じエキスパートを2回選ぶ場合、mul_mat_idは非対応

高い--moe-cache-cpu-shareでは空のキャッシュが埋まりにくく、必要以上のミスがCPUに回るため、開始条件を設けています。CPU部分はグラフ結果に蓄積し、llama_context::process_ubatchでグラフ入力設定後・計算前にキャッシュへ適用します。ggml-backendは次の分割の計算済み入力を現在の分割の実行前にホストメモリへコピーし、CPU分割がデバイス分割を待たずに走るようにしました。CPU部分の結果は非同期にコピーして戻します。

試し方: moe-cache-tune.py

moe-cache-tune.pyはllama-serverコマンドをラップして使います。

• 2種類のキャッシュサイズと3種類のCPUシェアでサーバーを起動
• マシンのアップロード速度とCPU速度を推定
• 実行時のルーティングをキャッシュロジックのコピー(ClassLRU方針を含む)で再生し、他のサイズとシェアを予測
• 推奨する--moe-cache-mibと--moe-cache-cpu-shareを出力

環境変数LLAMA_MOE_CACHE_TRACEを設定すると、各ルックアップで選ばれたエキスパートを書き出し、スクリプトのシミュレーターで利用できます。引数とスクリプトの説明はdocs/moe-cache.mdにあります。

開発体制

am17anがMoEキャッシュのコードオーナーに加わり、Gaurav Garg、Georgi Gerganov、Alexey Kopytkoが共同著者です。コミットには、Mac Studio(M5 Ultra、256GB)上のGLM-5.3-Flash-GGUFによる支援(Assisted-by)の記載があります。