KIV: RTX 4070(12GB VRAM)で1Mトークンのコンテキストウィンドウを実現、再学習なし、ドロップインでHuggingFaceキャッシュを置き換え—DynamicCacheを使う任意のモデルで動作 [P]

Reddit r/MachineLearning / 2026/4/13

💬 オピニオンDeveloper Stack & InfrastructureIdeas & Deep AnalysisTools & Practical Usage

要点

  • KIV(K-Indexed V Materialization)はミドルウェア層であり、HuggingFace transformersの標準KVキャッシュをティア(階層)付きメモリ方式に置き換える。最近のKVエントリはVRAMに保持し、古いものはシステムRAMへ移して、デコード時に取り出す。
  • この手法では、Kベクトルを構造化された検索可能なインデックスとして用い、各デコードステップごとにおよそトップ~256件の最も関連性の高いVエントリを選択する。これにより、デコード速度が総コンテキスト長にほぼ依存しないことを狙う。
  • 単一のRTX 4070(12GB VRAM)上でGemma 4(4-bit)を用いたベンチマークでは、KVのオーバーヘッド約12MBで1Mトークンまで対応でき、GPU全体の使用量は約6.5GBと報告されている。needle-in-haystack型および電話帳(phonebook)スタイルの検索テストでも良好な結果を示す。
  • 統合は「ドロップイン」でのHuggingFaceキャッシュ置換として説明されており、モデルの重みは変更せず、再学習/蒸留も不要で、DynamicCacheを使う任意のモデルで動作するはずである(Gemma 4、Qwen2.5、TinyLlama、Phi-3.5でテスト済み。MQA/GQA/MHAにも対応)。
  • 制限として、bounded prefillのもとでは見た目が似た密なデータに対して一部の損失が生じること、衝突の識別や二段推論(two-hop reasoning)で課題があること(一部は4-bitの2Bモデルが苦戦していることに起因するとされる)、またCPUからGPUへの転送がボトルネックとなりデコード速度の改善が制約されることが挙げられる。

しばらくこれに取り組んでいて、共有する準備ができたと思ったので投稿します。KIV(K-Indexed V Materialization)は、HuggingFace transformers の標準KVキャッシュを、階層型のリトリーバル(検索)システムに置き換えるミドルウェア層です。短く言うと、最近のトークンはVRAM上に正確なまま保持し、古いK/VはシステムRAMへ移し、さらにKベクトルを検索インデックスとして使って、デコード手順ごとに関連性の高いVエントリだけ(およそ~256件)を引き戻します。

Gemma 4 E2B(4-bit)を用いた、4070 12GBでの結果:

  • 1Mトークン、KIV VRAMオーバーヘッド12MB、総GPU使用量は約6.5GB
  • 1Mコンテキストで 1M tok/s(GPU時間で8〜10 tok/s)、4Kで12.9 tok/s
  • 4K〜32Kにわたるニードルインハイスタックテスト70/70通過
  • 58Kトークンでの完全な電話帳検索(固有名)
  • 1Mでのprefillは約4.3分(初回のコスト)
  • デコードはコンテキスト長にほぼ依存せず一定

これを成立させている核心の発見は、Kベクトルが滑らかで構造化されているため、検索インデックスとして非常に優れている点です。Vベクトルは高エントロピーでカオスなので、圧縮を試みるのではなく、必要に応じて取り出してください。Kを使って、どのVエントリがその時点でVRAM上に存在する価値があるかを判断します。

モデルの重みは一切変更しません。再学習や蒸留も行いません。HuggingFaceのキャッシュ・インターフェースにフックし、カスタムの注意(attention)関数を登録します。モデル側は、自分が階層型メモリシステムと通信していることをまったく認識しません。DynamicCacheを使う任意のモデルで動作します。Gemma 4、Qwen2.5、TinyLlama、Phi-3.5で、MQA/GQA/MHAにわたってテスト済みです。

実際の制限もあり、それらについてはリポジトリ内で正直に明記しています。境界付きprefillでは、密に似た見た目のデータに対していくらか情報が失われます。衝突の曖昧さ解消は機能しませんが、これはキャッシュの問題ではなく、4-bitの2Bモデルが苦労しているせいです。同様の理由で2ホップ推論にも失敗します。CPU RAMは線形に増えます(1Mで5.8GB)。

それでも引き続きデコード速度の最適化を進めています。特に長いコンテキストでの最適化です。現状のボトルネックは、モデルそのものではなく、取得したトークンのCPUからGPUへの転送です。ここには改善の余地がたっぷりあります。

GitHub: github.com/Babyhamsta/KIV(ローカルのpipパッケージとしてインストール可能。公式のpipパッケージはまだありません)

アーキテクチャや結果についての質問には喜んで答えます。より大きいモデルで、より多くのVRAMがある場合にどうなるかもぜひ見てみたいです。試してみたい方がいれば歓迎します。

submitted by /u/ThyGreatOof
[リンク] [コメント]