llama.cpp に TurboQuant を実装 ── 2bit・3bit の KV キャッシュ圧縮

ロングコンテキストは、計算の問題である前にメモリの問題です。モデルをロードした後、入力するトークン——そして生成するトークン——ごとに増えていくのがKV キャッシュ、つまりアテンション用に保存される key と value です。32K トークンではモデルの重みに匹敵し、128K では重みを大きく上回ります。コンシューマー向け GPU では、この上限が「使えるコンテキスト長で動く」か「メモリ不足で落ちる」かの分かれ目になります。

そこで私たちは llama.cpp をフォークし、Google Research が 2025 年に発表した量子化手法 TurboQuant をベースに、この KV キャッシュを圧縮する新しい方法を追加しました。フォークは github.com/openalchemy/llama.cpp にあり、2 つの新しいキャッシュ型——turbo3(値あたり 3 ビット)と turbo2(値あたり 2 ビット)——を提供します。既存の -ctk / -ctv フラグにそのまま渡せます。本記事では、その仕組みと、実際に何が得られるのかを解説します。

なぜ KV キャッシュはこれほど重いのか

各アテンション層は、トークンごとに key ベクトルと value ベクトルをキャッシュします。総サイズは 層数 × ヘッド数 × head_dim × コンテキスト長 × 2(K と V) × 値あたりバイト数 に比例します。モデルに手を入れずに動かせるのは最後の項だけです。FP16 で保存すると値あたり 2 バイト。標準の llama.cpp でも q8_0(約 1 バイト)や q4_0(約 0.5 バイト)に落とせますが、KV ベクトルには大きな外れ値があり、均一なグリッドは解像度の大半をその外れ値の追従に費やしてしまうため、素朴な低ビット量子化は品質をすぐに損ないます。

TurboQuant の何が違うのか

TurboQuant の着眼点は、量子化する前に回転させることです。各ベクトルにランダムな直交行列を掛けると、エネルギーが全座標へ拡散し、各座標の分布はほぼ同一で外れ値のないものになります。回転後は、座標ごとのシンプルなスカラー量子化器(最適な Lloyd–Max コードブック)が、与えられたビット予算に対してほぼ最良であることが理論的に保証されます。しかもデータ非依存——キャリブレーションセットも、モデルごとの調整も、学習も不要です。回転は読み出し時に元へ戻すので、アテンションはほとんど歪みのない往復を経たベクトルを見ます。

問題はコストです。密なランダム回転はホットパス上での行列積になります。私たちのエンジニアリング上の工夫は、密な Haar 行列を 高速ウォルシュ・アダマール変換(FWHT) に置き換えたことです。これは加算と減算だけで構成される直交変換で、O(d²) ではなく O(d log d) で動きます。期待値の上では同じエネルギー拡散の性質を持ち、自己逆変換である(デコードは同じカーネルをもう一度走らせるだけで、スケールは softmax に畳み込める)ため、浮動小数点の乗算が一切ありません。これにより、トークンごとの回転をキャッシュ書き込みパスの中に入れられるほど安価にできます。

どこまで小さくなるのか

各ブロックは 128 個の値と 1 つの FP16 ノルムをパックします。これは turbo3 で値あたり 0.39 バイト、turbo2 で値あたり 0.27 バイトになります——FP16 のおよそ 5 分の 1・7 分の 1 で、q4_0 よりも小さく、しかも生の 3 ビット・2 ビットグリッドよりはるかに品質を保つ回転を伴っています。

実モデルでの数値

値あたりの比率も結構ですが、重要なのは取り戻せる VRAM です。私たちは RTX 5080(16 GB) で 32,768 トークンのコンテキストにて、FP16 KV キャッシュと turbo3 の合計 VRAM を 2 つのモデルで実測しました。Qwen2.5-Coder-14B(Q4_K_M)では turbo3 が 4,418 MiB(約 4.3 GB)を解放し、Qwen3.5-9B(Q4_K_M、head_dim 256)では 828 MiB を解放しました。この 4.3 GB が、14B のコーディングモデルが本来なら載らないカードでリポジトリ丸ごと分のコンテキストを保持できるかどうかの差になります。

小さいことは遅いことを意味しません。圧縮されたキャッシュが演算ユニットの近くに常駐し続けるため、同じ実行で FP16 ベースラインに対しプロンプト処理 1.24 倍、生成 1.47 倍の速度を記録しました。回転カーネルのコストは十分に元が取れています。

この手法は 1 つのヘッド形状に縛られません。上記の 9B は head_dim 256(128 ブロックを 2 つ重ねたもの)、14B は 128 を使っています。head あたりの次元が 128 の倍数であるモデルなら対応します——Qwen 2.5 / 3 / 3.5 系、Llama 3.x、Mistral、Mixtral、Gemma が含まれます。

使い方

いちばん簡単なのは、コードを一切書かない方法です。TurboQuant はすでに OpenAlchemy Engine デスクトップアプリ(v0.5.1、v0.3.0 以降で利用可能)に搭載されています。Settings → Runtime → KV Cache Quantization を開き、key と value の両方のキャッシュを TurboQuant 3-bit に設定するだけ——再ビルドもフラグも不要で、次にモデルをロードした時から有効になります。

llama.cpp を直接動かしたい場合は、フラグ 2 つです。フォークをビルドして、キャッシュ型を turbo3 に向けます(Flash Attention が必須で、カーネルは現状 CUDA のみです):

./llama-cli -m Qwen2.5-Coder-14B-Instruct-Q4_K_M.gguf \
  -ngl 99 -fa 1 -c 32768 \
  -ctk turbo3 -ctv turbo3 \
  -p "def fibonacci(n):" -n 64

turbo2 も同じ方法で使え、極端なコンテキスト長のために最後の 1 ギガバイトまで絞り出したい時に適しています。2 ビットでは品質とのトレードオフが現実的でモデル依存なため、実験的・オプトインとして提供しています。既定で選ぶのは turbo3 です。私たちのテストでは、ロングコンテキスト用途で実質的に品質無劣化でした。

現在の実装と、これから

現在のパスは GPU 上でエンドツーエンドに動きます。キャッシュは圧縮して保存し、アテンションカーネルの直前で FP16 に逆量子化します。ロードマップの次は、逆量子化を Flash Attention のタイル内に融合し、key を一度も FP16 にマテリアライズしないようにすること、そして CUDA 以外のバックエンド対応です。CPU リファレンスパスは速度ではなく正しさのために存在します。

フォークはオープンソースです。コード・カーネル・テストハーネスは github.com/openalchemy/llama.cpp で読めます。何もビルドせずに結果だけ欲しい場合は、OpenAlchemy 推論プラットフォームを支えるエンジンの一部として利用できます。