しばらくこれに取り組んでいて、共有する準備ができたと思ったので投稿します。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がある場合にどうなるかもぜひ見てみたいです。試してみたい方がいれば歓迎します。
[リンク] [コメント]




