vToken、トークン単位の仮想化でKVキャッシュを回収——H100単体でスループットが最大37%向上

LLMサービングでは通常、PagedAttentionを用いてKVキャッシュを固定サイズのブロックに分割し、アロケータレベルの断片化を回避する。一方、H2OやScissorhandsなどの手法は、個々のトークン単位で保持または退避を判断する。同じブロック内に有効なトークンが1つでも残っていれば、ほかの空き領域をリソースプールへ返却できない。vTokenの予備評価によると、16Kコンテキストにトークン単位の退避を適用した場合、ブロック内部の無駄が40～60%に達することがある。

vTokenは、退避ポリシーと物理ブロックの間に仮想化レイヤーを追加する。リクエストごとのTokenTableには、論理トークン、生存状態、現在のブロックおよびoffsetが記録される。ポリシーは無効なトークンをマークするだけで、runtimeがメモリ圧力の上昇時に、使用率の低いブロック内で生存しているKVを新しい位置へ移動する。コピーはforward完了後に独立したCUDA streamで実行され、次のattention処理はCUDA eventで完了を待つ。このため、PagedAttention kernelの変更やCUDA Graphの再キャプチャーは不要だ。16Kシーケンスのマッピングmetadataは約256KBとなる。

著者らは、vLLM 0.18.0、PyTorch 2.10、80GBのH100単体という環境で、Mistral-7B、Llama-3.1-8Bに加え、容量評価としてQwen2.5-14Bをテストした。同じ退避判断を採用しながらブロックを圧縮しないNaive-Evictと比較して、vTokenは保持ブロック数を27.2～72.3%削減した。Mistral-7Bでは、SLA制約下のスループットが9.9～37.3%向上し、KV予算が制限された環境で実現可能な最大同時実行数は最大2倍になった。新しい退避ポリシーの統合に必要なコード量も、500行超から50行未満へ減少した。

ただし、モデル容量を無償で増やせるわけではない。データ移動は依然としてデコード処理とGPUリソースを奪い合い、トークン退避による品質低下についても上位レイヤーのポリシーが責任を負う。プロトタイプが対象とするのは単一ノード・単一GPUのfast pathのみであり、共有prefixブロックについては移動させない保守的な処理を採用している。今後は、マルチGPUのtensor parallelism、デバイス間のKV移動、混合ワークロードにおいて、圧縮による利得が同期および帯域幅のコストを相殺できるかを検証する必要がある。