KV Cache Offloadingとは？GPUからCPU・SSDへ退避する条件
LLM推論では、model weightだけでなくKV CacheもGPU memoryを消費する。長いpromptや多くの同時requestを扱うと、再利用できるprefixのCacheがあっても、空き容量を確保するためにGPUから追い出さなければならない。
KV Cache Offloadingは、再利用候補のKV blockをCPU memory、local SSD、remote storageなどへ移し、必要になったときGPUへ戻す仕組みである。GPU memoryの実効容量を広げられる一方、転送が遅ければpromptをもう一度Prefillする方が速い場合もある。
退避するのはmodel weightではなく計算済みKVである
移すのは、過去tokenをAttentionで参照するために計算したKeyとValueである。model本体をCPUへ移すweight offloadとは目的が異なる。また、実行中requestを継続するための一時的なswapと、完了済みrequestの共通prefixを後続request向けに保存する処理も区別した方がよい。
GPU、CPU、SSDを階層として扱う
保存先は容量が大きくなるほど一般に遅くなる。
| 階層 | 特徴 | 主な制約 |
|---|---|---|
| GPU HBM | Attentionから直接使える | 容量が小さく高価 |
| CPU DRAM | GPUより大容量 | PCIeやNVLinkの転送時間と帯域 |
| local NVMe SSD | さらに大容量 | read latency、write量、SSD寿命 |
| remote storage | 複数workerから共有できる | network latency、帯域、障害範囲 |
LMCacheはGPU、CPU、local disk、remote backendへKV Cacheを配置し、推論engineをまたいで再利用する設計を示している。NVIDIA DynamoのKV Block ManagerもGPU、CPU、diskを階層として扱う。容量だけを見てSSDへ大量保存すると、書き込みと再読込がbottleneckになるため、各階層へ何を残すかというadmission policyが必要である。
再読込と再計算のどちらが速いか
offloadしたKVを再利用する価値は、次の関係で判断できる。
[ T_{lookup} + T_{load} < T_{prefill_recompute} ]
左辺はCacheの検索とGPUへの読込時間、右辺は同じprefixをもう一度Prefillする時間である。長いprefixではoffloadが有利になりやすい。短いprefix、高速なGPU、混雑したPCIeやSSDでは再計算の方が速いことがある。
判断にはprefixが再利用される確率、KV blockのbyte数、GPUの混雑、転送帯域、modelとdtypeの互換性、SSD write量も含める。NVIDIA Dynamoのdisk offloadには、同じblockが一定回数以上現れた場合だけSSDへ書くfilterがある。一回しか使われないpromptまで永続層へ送らない考え方である。
保存と読込を生成処理に重ねる
GPUからCPUへのcopyを同期的に待てば、その間に次token生成が止まりうる。実装ではblockをまとめて転送し、GPU計算とdata movementを重ね、backgroundで次の階層へ流す。
読込側も同様である。request受付後にSSDを読み始めるだけではTTFTが伸びる。routing時点でCacheの所在が分かれば、queue待ちの間にprefetchできる。ただし、使われないblockを先読みすれば帯域を浪費する。
CPU DRAMはtokenizer、request queue、vector search、OS page cacheとも共有される。offload上限を設け、hostのmemory pressureとswap発生を監視する。
Cacheがあるworkerへ送るか、Cacheを移すか
複数GPU workerでは、同じprefixのKVがworker Aにあり、空いているworker Bにはない状況が起こる。Aへroutingする方法と、BへKVを転送する方法がある。
Aのqueueが長ければ、local CacheがあってもBで再計算する方が早い。routerはhitの有無だけでなく、queue時間、transfer時間、Prefill再計算時間を比較する必要がある。
Prefill・Decode分離で新しいKVを別workerへ渡す経路も、同じnetwork帯域を使う場合がある。offloadと同時に負荷試験する。
Cache keyに実行条件を含める
同じ文字列でも常に同じKVを再利用できるとは限らない。少なくともmodel IDとweight revision、tokenizerとtoken ID列、KV dtype、parallel構成、LoRA、multimodal入力、tenant、権限、cache saltをkeyやnamespaceへ反映する。
rolling deploymentではversionごとにnamespaceを分け、旧versionのCacheを期限付きで捨てる。共有Cacheにはprompt由来の情報が含まれるため、別tenantから存在確認やtiming差を観測できないようにする。暗号化、access control、保持期間、削除手順も必要である。
Offloadは正しさのための保存先ではない
KV Cacheは再生成できる派生dataである。CPU processの再起動、SSD故障、remote backend停止、checksum不一致が起きたら、missとしてPrefillを再実行できる設計にする。metadataだけが残る場合や部分blockしか読めない場合も安全にmissへ倒す。
ただし、backend障害中に全requestが再計算へ切り替わればGPU負荷が急増する。Concurrency Limit、queue、load sheddingと組み合わせ、回復時も一斉にCacheを温めない。
監視するmetric
hit率だけでは効果を判断できない。短いprefixのhitが多くても転送費用の方が大きい可能性がある。
- 階層ごとの保存block数、byte数、使用率
- D2H、H2D、disk read・writeのblock数とbyte数
- matched token数と再計算を省けたtoken数
- 階層ごとのhit率、lookup latency、load latency
- offload queue、prefetch queue、転送帯域
- Cache hit時とmiss時のTTFT
- 再計算へ切り替えた回数、version mismatch、checksum error
- SSD write量、容量、寿命指標
- host memory pressureとOS swap
最終的にはoffloadなしと比較したTTFT、GPU利用率、throughput、1 request当たり費用を見る。
負荷試験
- 長い共通prefixを高頻度で再利用する
- すべて異なるpromptで再利用がない
- 短いprefixと長いprefixを混在させる
- CPU、SSD、networkの帯域を圧迫する
- GPU workerを増減し、Cache localityを変える
- modelとLoRAのversionを切り替える
- offload backendを停止し、再計算へ戻れるか確認する
- metadataだけが残る状態と破損blockを注入する
- cold start後に大量requestを送る
- Prefix Cache、Prefill・Decode分離、offloadを同時に有効化する
平均だけでなくp95、p99のTTFT、transfer queue、recompute量を見る。SSD容量を埋める長時間試験も必要である。
関連技術との違い
- PagedAttentionはGPU上のKVをblock単位で配置する
- Prefix Cachingは再利用可能な共通prefixを特定する
- KV Cache OffloadingはblockをGPU外の階層へ移す
- Prefill・Decode分離は推論の二段階を別worker poolで処理する
- Cache-aware routingはKVがあるworkerを選ぶ
これらは組み合わせる技術である。ただし組み合わせるほどmetadata、転送経路、failure modeが増える。まずGPU内のPrefix Cacheで効果を測り、容量不足と再利用量が確認できた段階でCPU、SSD、remote層へ広げるのが安全である。
Offloadの目的は再計算の削減である
KV Cache Offloadingは、GPUから追い出した計算済みKVをCPUやSSDへ残し、後続requestのPrefillを省く。ただし、保存すれば必ず速くなるわけではない。再利用されるprefixを選び、検索と転送が再計算より短い範囲で使う必要がある。
設計の中心は保存容量ではなく、再計算時間、転送帯域、Cache locality、versionとtenantの境界、missへ安全に戻る経路である。本番promptの長さと再利用分布で損益分岐を測る。