Monarch v3: NESに着想を得たKVページングでLLM推論を78%高速化

TL;DR: トランスフォーマー向けにNESに着想を得たメモリページングを実装しました。1.1Bパラメータのモデルでは、推論が78%高速（17.01 → 30.42 tok/sec）になり、ほぼゼロのVRAMオーバーヘッドです。アルゴリズムはオープンソースで、完全にベンチマーク済みで、すぐに使える状態になっています。

課題

KVキャッシュはシーケンス長に対して線形に増加します。4Kトークンになると、その大部分は未使用のままになります—古いトークンよりも最近のトークンのほうがはるかに重要なのに、それでもすべてをVRAM上にフル精度で保持し続けています。

標準的なアプローチ（量子化、プルーニング、蒸留）は侵襲的です。私たちはもっと単純な方法を望みました。つまり、古いものを邪魔にならない場所へ移すだけです。

解決策: NESに着想を得たページング

ゲームボーイのメモリバンキングシステムだと思ってください。キャッシュはホット領域（最近のトークン、フル精度）とコールド領域（より古いトークン、圧縮）の2つに分割します。新しいトークンが到着するたびに、古いトークンはホットストレージから追い出され、コールドストレージへ圧縮されます。トークンが（高い注意重みによって）昇格されると、再びホットへ移動します。

重要なトレードオフ: ホットウィンドウに対してのみフルの注意計算を行います。コールドトークンは、明示的な昇格があったときだけアクセスされます。これは標準的な注意機構とは本質的に異なり、「最近のトークンが支配的である」と仮定します。これは多くのタスクでは当てはまりますが、すべてではありません。

4つの要素が連携します:

1. ウィンドウ付き注意（高速化エンジン）
   • 注意はホットウィンドウのみ（デフォルトで約512トークン）
   • 古いトークンでも、アクセスされれば昇格できる
   • 前提: 近接性は注意の強いシグナル
   • 未検証: 基準（ベースライン）に対するフル生成品質への影響
2. TurboQuant圧縮（コールドKVのサイズを約97%削減）
   • コールドKVを4ビット整数へ量子化
   • 類似度のために極座標符号化（半径 + 偏角ビン）
   • 残差補正（値あたり1ビット）
   • アクセス時に最小オーバーヘッドでデコード
3. スライディングウィンドウの追い出し
   • 直近Nトークンはデフォルトでホットのまま
   • 古いトークンはコールドストレージへ圧縮
   • 「重要な」トークンを事前に知る必要はない
4. 注意重みによる昇格
   • 高い注意を受けたトークンは再びホットへ移動できる
   • スティッキー機構がスラッシングを防ぐ
   • 閾値ベースで、誤った昇格を避ける

ベンチマーク結果

セットアップ: TinyLlama-1.1B fp16、生成50トークン、ウィンドウ付き注意を有効化

モード	スループット	VRAM	ホットウィンドウ
標準（フル注意）	17.01 tok/s	2112 MB	—
Monarch-v3（ウィンドウ付き）	30.42 tok/s	2131 MB	512トークン
増加	+78.7%	+0.9%	—

この大きな高速化は、注意計算を最近のトークンに対してのみ行うことに起因します。圧縮によって少しVRAMが節約されますが、主な勝ち筋ではありません。

重要な注意点: このベンチマークは生成品質ではなくスループットを測定しています。ウィンドウ付き注意 + 昇格が、フル注意と区別できないテキストを生み出すかどうかは検証していません。近接性の仮定は多くのタスクでうまく機能する一方、検索中心のクエリや文脈依存のクエリでは失敗する可能性があります。

仕組み（簡略化したデコードループ）

for step in 1..100: q = project_query(next_token) # Standard: compute attention over ALL cached tokens # Monarch: compute attention only over HOT window scores_hot = q @ kv_hot.T # ~512 tokens instead of 4096+ # Optional: Check if cold tokens should be promoted # (only if attention scores suggest they matter) if promotion_enabled and max(scores_hot) < promotion_threshold: kv_cold_promoted = decompress(cold_pages) scores_cold = q @ kv_cold_promoted.T if max(scores_cold) > threshold: promote_cold_to_hot() # Softmax over [hot + promoted], apply attention # Old tokens fall out of hot window if len(kv_hot) > window_size: compress_to_cold() 

高速化の要点は、ほとんどの古いトークンについて注意計算をスキップすることです。これによって生成品質が保持されるかどうかは、まだ未解決の問いです。

現在の状況

実装: custom cache backendを備えたHugging Face Transformersで開発中
ベンチマーク: 複数のシーケンス長に対する完全な検証
オープンソース: Apache 2.0、フォーク可能な状態
論文: 完全な技術仕様（NESに着想を得たページング、圧縮方式、評価手法）

次: コールドデコンプレッション向けのCUDAカーネル融合（さらに効果を押し上げます）

試してみる

クローンして実行:

git clone https://github.com/JohannaWeb/Monarch.git cd Monarch # 依存関係をインストール pip install -r requirements.txt # Project Falconの知識でTinyLlamaを学習 python train_tinyllama_fp16.py # 標準推論とページング推論をベンチマーク python src/benchmark_monarch.py 
d --model models/tinyllama_fp16 
 --mode both 
 --max-new-tokens 100 
 --promotion-threshold 0.15 
 --sticky-threshold 3 
 --json 

わかっていること／わかっていないこと

検証済み:

• スループットの改善（短いシーケンスで+78.7%）
• VRAMオーバーヘッドは最小（+0.9%）
• 実装は安定しておりクラッシュしない

想定はしているが検証していない:

• ウィンドウ付き注意で生成品質が保持される
• 近接性仮説が多様なタスクで成り立つ
• 効果がより長いシーケンスやより大きいモデルへ転移する
• 昇格機構が重要なコールドトークンを正しく識別する

未実装:

• ベースラインに対する完全なBLEU/パープレキシティ評価
• より長いシーケンスのベンチマーク（>1000トークン）
• 検索中心のタスクにおける品質評価
• マルチトークンのバッチデコード（単一シーケンスのみ）

FAQ

Q: ウィンドウ付き注意は生成品質を劣化させますか？
A: 不明です。私たちは出力品質ではなくスループットとVRAMをベンチマークしました。近接性仮説はもっともらしいです（最近の文脈が最も重要だ）が、ベースラインとのBLEU/パープレキシティのベンチマークは実行していません。検証上の実際のギャップです。

Q: KVキャッシュの量子化に関する論文はどうなりますか？
A: 私たちはコールドトークンを量子化し、ホットトークンは量子化しません。ホットトークンはフル精度のままです。ただし、主な高速化は圧縮ではなくウィンドウ付き注意によるものです。

Q: これはどんなタスクに向いていますか？
A: おそらく: チャット、要約、RAGのように最近の文脈が支配的なケースです。可能性が低い: ニードル・イン・ヘイスタック型の検索や、古いトークンが重要になるメモリ重視のタスクです。

Q: バッチ推論はどうなりますか？
A: 現在の実装は単一シーケンスです。バッチ化にはページ管理を慎重に行う必要があります（今後の課題として残しています）。

Q: vLLMやSGLangと一緒に使えますか？
A: まだできません。標準的なトランスフォーマーに対する概念実証（proof-of-concept）です。統合するには、これらのシステムがcustom cache backendを採用する必要があります。

Claude（AIペアプログラミング）とJohannaによって作成

リポジトリ: https://github.com/JohannaWeb/Monarch
論文: リポジトリ内のmonarch_nes_paper.htmlを参照してください