vLLMの階層型KVキャッシュ、GPU外への拡張とノード間共有を実現

長いコンテキストや多ターン対話では、KVキャッシュがアクセラレーターのメモリ容量をすぐに圧迫します。古いデータが追い出された後に同じ内容が必要になると、従来は最初から再計算する必要がありました。vLLMのTiered KV Cache Offloadingは、追い出されたKVチャンクを下位階層に保持し、必要なときに読み戻す仕組みを提供します。

CPUメモリを中心にした経路

この設計の特徴は、GPUから外部ストレージへ直接アクセスさせず、すべてのKVデータをホストメモリに通すことです。オフロード時は、まずアクセラレーターからCPU DRAMへコピーします。このコピーが終わればGPUメモリを解放でき、その後のストレージ書き込みやリモート転送はホスト側のコピーから非同期に進みます。再ロード時は逆に、二次階層からホストメモリへ戻してからGPUへ投入します。そのため、GPUメモリはデータが実際に利用可能になるまで長時間予約されません。

ホストメモリは単なる中継バッファではなく、LRU/ARCを備えた一次キャッシュです。スケジューラーはまずこの層を確認し、ミスした場合に設定順で二次階層を検索します。要求内の異なるチャンクが、ホスト、ファイルシステム、リモートノードなど別々の階層から供給される場合もあります。

異なる構成をまたぐ共通形式

テンソル並列では、各アクセラレーターがKVキャッシュの一部だけを保持します。vLLMはそれらの分片を共有ホスト領域へ集約し、1つのページに1レイヤーの1ブロックを置き、複数ランクのKVヘッドを連続領域にまとめる標準レイアウトを使います。これにより、二次階層に対する細かな転送を減らし、より大きなI/Oとして扱えます。

ホスト側の形式が固定されているため、アクセラレーターの種類、注意機構のバックエンド、並列度が異なるノード同士でも、GPU側の配置を変換せずにKVデータを共有できます。二次階層はアクセラレーター固有のAPIを理解する必要がなく、一般的なCPU向けI/Oで実装できます。

二次階層と利用場面

• ファイルシステム：内容アドレス方式でチャンクを保存し、同じマウント先を使うインスタンス間で共有できます。
• オブジェクトストレージ：S3互換サービスを利用し、ネットワーク越しの共有を比較的低コストで実現します。
• P2P転送：ZMQで調整し、NIXL経由のRDMAでホスト間のデータを転送します。

P2PはPrefill/Decode分離と負荷分散に適しています。Prefill側がホストメモリに用意したチャンクをDecode側が取得でき、チャンク単位のPrefillでは計算と転送を早い段階から重ねられます。負荷の高いインスタンスから余裕のあるインスタンスへKVを移す用途も想定されています。GPUごとの細かな転送をホスト側で集約してから送るため、ネットワークI/Oの効率化も期待できます。

意義と注意点

この仕組みは、KVキャッシュを単一GPU内の一時データから、階層化・永続化・共有可能な推論基盤へ変えるものです。再計算を減らし、クラスターの実効的な提供容量を高める可能性があります。一方で、効果はホストメモリ容量、ストレージ遅延、帯域幅、キャッシュヒット率に左右されます。遠隔階層は再計算を避けても、転送コストそのものをなくすわけではありません。それでも、複数ノードにコンテキスト状態を広げるための実装上の道筋を、vLLMに与える設計だと言えます。

vLLM Blog