Fathom、クエリごとにKVキャッシュの読み取り深度を最適化

長文脈推論では、ボトルネックがモデルの計算量ではなくメモリ転送になる場合があります。エージェントのセッションが百万トークン規模に達し、多数のセッションが同時に常駐すると、KVキャッシュと候補順位付け用のインデックスはホストRAMへ移されます。デコードのたびに全キーを走査してTop-k候補を選ぶ必要があるため、最終的に使うKV行を読み出す処理よりも、インデックスの走査が重くなります。

Fathomの特徴は、すべてのクエリに同じ走査精度を適用しないことです。各クエリが、キーの各チャネルから何ビットを読むかを決めます。4ビットのKキャッシュはチャネル単位のビットプレーンとして保存され、あるチャネルの先頭tプレーンを読むことは、そのチャネルをtビット量子化値で扱うことと同じです。そのため、部分読み出しでも連続したメモリアクセスを保ち、段階的に精度を高められます。

総ビット予算の範囲内で、Fathomは現在のクエリにとって重要なチャネルへ多くのビットを配分します。この逆ウォーターフィリングによる割り当ては閉形式で計算でき、単一のTritonカーネルで実行されます。

論文が示す主な結果は次の通りです。

• A100上のQwen3-8B、百万トークンの文脈で、Double Sparsity、Loki、SparQ系の136ビット走査と比べ、GPU時間ベースで1.67倍高速でした。
• SparQの68ビット読み出しと同じGPU時間で、Fathomは読み取るバイト数を18%削減し、7つのモデル・文脈設定のうち6つでより低い注意誤差を示しました。
• RULER系タスクでは、各トークンの走査結果が厳密なTop-kデコードと一致しました。実際のOpenHandsコーディングエージェントでは、92ビットで最も精度の高い136ビット走査と同等のステップ一致率に達しました。

この研究の意義は、インデックスを一律に圧縮するのではなく、クエリごとに読み出し予算を配る点にあります。KVデータがホストメモリにある推論基盤では、不要な転送を減らすことが算術演算の削減以上に重要になり得ます。一方で、インデックスがGPUメモリに常駐する場合、走査は演算律速となるため、ビットを減らしても乗加算処理は残り、高速化にはつながりません。論文は校正データの違いが結果へ与える影響も小さいと報告し、コードや結果ファイル、実験チェーンを公開しています。

出典：Hugging Face Daily Papers