2.78兆パラメータのLLMをCPU 1基・8GB RAMで実行:Kimi K3 Cエンジンの内部

176KBのC99バイナリが、ストリーミング、MXFP4、巧妙なメモリ管理により、わずか8.24GBのRAMで2.78兆パラメータのKimi K3モデルを実行する方法をご覧ください。

ノートパソコンで2.78兆パラメータの言語モデルを実行することを想像してみてください。量子化や蒸留、その他の妥協を一切行わず、完全なオリジナルのKimi K3、1.56TBの重みを持ち、データセンターのGPUクラスタで得られる出力とバイト単位で同一の出力を生成します。それがまさにkimi-k3-in-cプロジェクトが達成していることであり、システムエンジニアリングの傑作です。

これはおもちゃのデモではありません。ポータブルなC99で書かれた完全に機能する推論エンジンであり、BLASもディープラーニングフレームワークもGPUサポートもありません。エンジン全体は176KBのバイナリにコンパイルされます。秘密は、モデルアーキテクチャの活用、アグレッシブなメモリストリーミング、そしてすべてのバイトがどこにあるかについての深い理解の組み合わせです。

問題:収まらないモデル

Kimi K3には2.78兆のパラメータがあり、bfloat16精度では5.56TBのメモリが必要になります。リリースされたチェックポイントは1.56TBで、それでもコンシューマーマシンには遠く及びません。しかし、Kimi K3はMixture-of-Experts(MoE)モデルです。各トークンに対して、レイヤーごとに896エキスパートのうち16だけがアクティブになります。これは約1,040億のアクティブパラメータ、つまり全体のわずか3.7%です。

重要な洞察:他の96.3%のパラメータはRAMに存在する必要はありません。到達可能であればよいのです。非アクティブなエキスパートをディスクからオンデマンドでストリーミングすることで、メモリ要件は劇的に削減されます。

4つの削減

このプロジェクトは、理論上のbfloat16ベースラインから4つの主要なステップで675倍のメモリ削減を達成しています。

  1. MXFP4量子化:ルーティングされたエキスパートはすでにMXFP4形式(マイクロスケーリング4ビット浮動小数点形式)で保存されています。各重みは4ビットのニブルで、32重みごとに共有の8ビット指数があります。これにより、エキスパートのストレージが5.45TBから1.447TBに削減されます。

  2. スパーシティ:トークンごとに896エキスパートのうち16だけが発火するため、1.447TBのエキスパートを完全に常駐させる必要はありません。必要に応じてディスクからストリーミングされます。

  3. KDAアテンション:93レイヤーのうち69レイヤーがKimi Delta Attentionを使用します。これは固定サイズの状態を持つリカレントアテンションメカニズムで、コンテキスト長に応じて成長しません。これにより、ほとんどのレイヤーでKVキャッシュの爆発が排除されます。

  4. MLA圧縮:残りの24レイヤーはMulti-head Latent Attentionを使用し、位置ごとにヘッドごとのキー/バリューペアの代わりに単一の576次元の潜在ベクトルをキャッシュします。これにより、KVキャッシュサイズが53倍削減されます。

これらの削減後、常駐セットは113.49GB(高密度トランク、埋め込み、出力ヘッド)だけになります。しかし、それでもノートパソコンには大きすぎます。最終的な削減:トランク自体をストリーミングします。

トランクのストリーミング:ダイヤルであって床ではない

トランク(93の高密度レイヤー)は108.81GBです。すべてのレイヤーがすべてのトークンで使用されるため、活用できるスパーシティはありません。解決策は、予算が許す限り多くのレイヤーをRAMに固定し、残りを単一のリングバッファを介してディスクからストリーミングすることです。

エンジンは、トークンごとにレイヤー0〜92を順に歩きます。これは循環スキャンであり、LRUキャッシュにとって最悪のケースです。そのため、キャッシュの代わりに固定プレフィックスを使用します:最初のNレイヤーは永続的に常駐し、残りはリングスロットを循環します。これにより、決定的なヒット率N/93が得られます。

例えば、--preset laptop(3GBトランク予算)では、10レイヤーだけが固定され、残りはディスクからストリーミングされます。結果:ピークRSSは8.24GBですが、トークンあたり26.5秒かかります。--preset server(110GBトランク予算)では、90レイヤーが固定され、速度はトークンあたり5.6秒に向上します。

エキスパートキャッシュ:測定の教訓

エンジンには、ルーティングされたエキスパート用のLRUキャッシュも含まれています。しかし、驚くべきことに、キャッシュは小さいサイズではほとんど役に立ちません。プロジェクトの測定によると、キャッシュを28スロットから1,344スロット(48倍の増加)に増やしても、トークンあたりの読み取りバイト数はまったく変わりません。

なぜでしょうか?Kimi K3はQuantile Balancingと呼ばれる技術を使用しており、エキスパートの使用状況をプール全体で平坦化します。ホットなエキスパートのサブセットがないため、LRUキャッシュは有用なものを保持しません。キャッシュがワーキングセットのかなりの割合(約36GBのアリーナ)を保持できるほど大きくなって初めて、キャッシュは役立ち始めます。

これは直感に反する最適化につながります:メモリをエキスパートキャッシュではなく、まずトランクに割り当てることです。固定128GBの予算では、トランクに110GB、キャッシュに13GBを割り当てると、逆の分割よりも1.69倍高速ですが、後者の方がキャッシュヒット率は高くなります。

ビット単位の再現性

最も印象的な側面の1つは、ビット単位の再現性への取り組みです。エンジンは、スカラー、OpenMP、AVX2のコードパスがすべて同一の結果を生成することを保証します。これは以下によって達成されます:

  • -ffp-contract=offを使用してFMA縮約を防ぎ、丸めを変更しないようにします。
  • 固定の合計順序で倍精度で累積します。
  • bf16を単純なシフト(ロスレス)でfp32に拡張します。

これにより、8GBのノートパソコンで実行しても224GBのワークステーションで実行しても、出力はバイト単位で同一になります。唯一の違いは速度です。

検証:動作の証明

プロジェクトには厳格な検証スイートが含まれています:

  • ウェイトレステスト:フルモデルと同じテンソルグラフを持つ13レイヤーのオラクルモデルを、PyTorchリファレンスと照合します。すべてのゲートが正確に合格します。
  • フルチェックポイント適合:全93レイヤーをPyTorchと照合し、最悪のエラーは許容誤差の0.00倍です。
  • ロジットパリティ:CエンジンのロジットはPyTorchリファレンスと要素ごとに一致し、最大差は7.87e-6です。
  • メモリラダー:8GBから224GBまでの12の異なるメモリ予算で、すべて同一のトークンIDを生成します。

結論

このプロジェクトは、慎重なシステムエンジニアリングで何が可能かを示す証です。1兆パラメータのモデルがデータセンターを必要としないことを証明しています。必要なのは、バイトがどこにあるか、そしてそれらを効率的に移動する方法についての賢い理解だけです。

MoEアーキテクチャ、メモリ効率の高い推論に興味がある場合、または単に印象的なCコードを見たい場合、kimi-k3-in-cは深く掘り下げる価値があります。READMEだけでも技術的な詳細の宝庫であり、コードはクリーンでよく文書化されています。

自分で試してみたい場合は、約1.7TBの空きディスク容量と、ダウンロードに多くの忍耐が必要です。しかし、その見返りは、おそらくすでに所有しているハードウェアで最大級のオープンモデルの1つを実行することです。

ソース

FareedKhan-dev/kimi-k3-in-c: 8.24 GBのRAMで単一CPU上で推論を実行する2.78兆パラメータのKimi K3。ポータブルC99:BLASなし、フレームワークなし、GPUなし。