ホームへ戻る

模型推論與執行期

llama.cpp、Kimi K3向けに循環状態ロールバックを追加――投機的デコーディングでホスト側チェックポイントの複製が不要に

新しいプレリリース版では、候補位置ごとにKDA状態と畳み込みウィンドウを保存し、拒否されたドラフトtokenを正しくロールバックできる。初期の単一マシンテストでは、低並列の推論で高速化が見られた一方、状態メモリはロールバック深度とサービスslot数に応じて大幅に増加する。

The GGML authors · Public domain · Image source
zh-Hant

llama.cppのプレリリース版b10853は、Kimi K3向けに「境界付き循環状態ロールバック」を有効化し、Transformer/循環ハイブリッドアーキテクチャで投機的デコーディングを使用する際の重要な欠落を補った。投機的デコーディングでは、まず小規模なドラフトモデルが複数のtokenを提案し、その後、メインモデルが一括検証する。メインモデルがドラフトの一部を拒否した場合、ランタイムは最後に受理された位置までキャッシュを復元しなければならない。従来のTransformerでは主にKV cacheを巻き戻すが、Kimi K3の69層のKimi Delta Attentionは、固定サイズの循環状態とQ/K/Vの短い畳み込みウィンドウも保持している。最終状態だけを保存すると、ドラフトが拒否された後に存在しないスナップショットを復元することになり、以降の計算が正しいシーケンスから逸脱する。

マージされた実装では、ロールバック可能な各位置について畳み込みウィンドウを保存し、llama.cppのrecurrent-attentionヘルパーを再利用してKDA状態のスナップショットを作成する。テストでは、チェックポイントの復元、複数シーケンスを分割したリプレイ、シーケンス間の分離を網羅しており、macOS CPUとNVIDIA Vulkanを対象としたテストはいずれも合格した。特に重要なのが非ゼロ値によるパディングテストだ。Kimi K3を許可リストに追加しただけでは、ゼロ初期化されたキャッシュによって不具合が覆い隠される可能性があるが、`0x3e`で埋めることで許容誤差を超えるlogitの差異を検出できた。

開発者は、量子化されたKimi K3、4つのサービスslot、7つのドラフトtoken、8基のRTX PRO 6000 Blackwellを使用してテストした。その結果、ロールバック方式はホスト側チェックポイント方式と比べ、文章書き換え、知識、コーディングの各ワークロードで約13~53%高速だった。ただし、これはplacement-controlled A/Bテストではなく、モデル構成と利用可能な空き容量が異なる。投機的デコーディングをまったく使用しない場合と比べると、並列度1ではTPOTが改善したものの、並列度4では逆に悪化した。コストも明確で、4 slot、7つのロールバック位置という条件では、循環状態キャッシュが約1.77 GiBから14.18 GiBに増加した。デプロイ担当者は次の段階として、単一リクエストの速度だけを見るのではなく、ドラフト受理率、並列度、追加の状態メモリを測定すべきだ。

出典

  1. llama.cpp b10853 release
  2. Kimi K3 recurrent-state rollback implementation and measurements
  3. Kimi K3 model card