KVキャッシュがGPUに収まらないとき: 循環アクセスとLRU、そしてCPUへの退避
一言でいうと
GPUのKVキャッシュの容量より多くの文書を交互に読み込ませると、最も長く使われていないものから追い出すLRUのせいで、同じ順序でもう一度読んだときにヒットが0になります。追い出されるKVをCPUメモリに退避しておけば(LMCacheの考え方)、再計算せずに読み戻せます。この測定では、最初のトークンまでの時間が576 msから48.6 msに短縮されました(約11.9倍)。代償は、最初に読むときに約4%遅くなることです。
なぜ必要なのか
vLLMは、同じ先頭部分(プレフィックス)を持つリクエストのKVキャッシュをGPUに置いて再利用します(prefix caching)。長い文書を先頭に付けて質問だけを変えるリクエストなら、文書部分の計算を省けます。問題はGPUメモリが小さいことです。この測定では、1.5Bモデルを載せたあとに残ったKVキャッシュの容量は21,520トークンだけで(サーバーログのGPU KV cache size)、文書1つが5,011トークンなので、4つ強しか入りません。文書が10個あると合計は50,110トークンで、容量の約2.3倍です。
容量が足りなければ何かを追い出す必要があり、このキャッシュは最も長く使われていないものから追い出します(LRU)。最近使ったものがすぐにまた使われると仮定する規則ですが、文書0から文書9まで順に読み、同じ順序でまた文書0から読む循環アクセスは、その仮定がちょうど逆になる場合です。最初のパスが終わった時点で残っているのは、最も最近読んだ文書6–9です。2回目のパスでは文書0がすでにないので、新しく計算して入れますが、そうすると最も古い文書6が追い出されます。文書1を入れると文書7が追い出されます。再び読む順序が追い出された順序と同じなので、最も古いものが、最も早くまた必要になるものになります。容量が少し足りないだけでもヒットが0になる理由です。実測でも、vLLMのprefix_cache_hitsは2回目のパスで0でした。
どう動くのか
LMCacheは、追い出されるKVをCPUメモリ(ディスクやRedisも可能)に退避しておき、同じプレフィックスが再び来たら、計算する代わりに読み戻します。このラボのシミュレーターは、その流れを3つのルールに絞りました。
- GPUでヒットしたら、そこで終わりです。CPUには触れません。
- GPUでミスしたがCPUにあれば、計算せずにCPUからGPUへ読み戻します。
- どちらにもなければ計算し、GPUとCPUの両方に保存します。
ルール1は測定ログが裏づけています。逆順に読む3回目のパスで、vLLMがGPUで解決したプレフィックスは21,488トークンで、LMCacheに問い合わせたトークンは28,755個でした。全体の50,243からGPUでのヒット分を引いた値です。ルール3は実際の動作を簡略化したものです。レポートは、最初の読み込みが約4%遅くなったことを、KVをCPUに退避するコストと解釈しています。また、文書10個(5万トークン)のKVが約1.3GBで、設定した上限1.5GBがそれを収められたと記しています。CPU側の容量にも上限があるため、上限より多くの文書を読み込ませると、CPUでも古いものから捨てられます。
実測結果
文書10個を、最初の読み込み、同じ順序での再読み込み、逆順の読み込みの3つのパスで読み込ませました。リクエストは1つずつ送り、最初のトークンまでの時間(TTFT、ms)を測りました。値はp50と平均です。
| 設定 | パス | p50 | 平均 |
|---|---|---|---|
| vLLMのみ | 最初 | 575.5 | 581.2 |
| vLLMのみ | 同じ順序で再度 | 576.2 | 575.9 |
| vLLMのみ | 逆順 | 505.4 | 343.6 |
| LMCache | 最初 | 597.8 | 605.4 |
| LMCache | 同じ順序で再度 | 48.6 | 49.6 |
| LMCache | 逆順 | 47.2 | 40.9 |
同じ順序で再度読むとき、p50は576から48.6 msへと約11.9倍、平均は約11.6倍速くなりました。パス1つを終えるまでの時間は、6.2秒から0.9秒でした。最初に読むときは、p50が575.5から597.8 msへと約4%長くなりました。vLLMのみの2回目のパスは、最初のパスと実質的に同じです(576対575)。キャッシュがまったく役に立たなかったという意味です。
LMCacheのログでは、2回目のパスのリクエスト1つが、約5,000トークンのうち4,864トークン(256の倍数19個)をCPUから読み戻し、残りの約150トークンだけを新しく計算しました。読み戻しには8–9 msかかり、同じ分量を新しく計算するには約560 msかかります。読み戻しが計算より60倍以上安いということが、この差のすべてです。
逆順に読むパスも、シミュレーターは当てます。vLLMのみの場合、直前に読んだ4つの文書(文書9、8、7、6。番号は測定ファイルのdoc値です)は、約30 msでGPUからヒットし、残りは約576 msでした。LMCacheを有効にすると、同じ4つの文書は約30 ms、残りの6つは約46–49 msでCPUから読み戻されます。ただし、vLLMのみの側の文書5は438 msで、ヒットとミスの中間の値になりました。文書を丸ごと出し入れするシミュレーターは、このような部分的なヒットを再現できません。
この数字をどこまで信じるか
- このシナリオは、LMCacheに最も有利です。同じ文書を再度読み、プレフィックスが5千トークンと長く、GPUキャッシュが足りません。利得は、再利用される長いプレフィックスの割合に左右されるとレポートは解釈していますが、今回はその割合を変えて測ってはいません。
- リクエストは1つずつ送りました(同時実行数1)。リクエストが重なるときの様子は、この測定では示されません。
- モデルは1.5Bと小さく、GPUは1枚です。大きいモデルでは差がさらに広がると予想されますが、この測定で確認したものではありません。
- 文書は乱数の単語で作った合成文書なので、実際の文書のようにプレフィックスが部分的に重なることはありません。
- 一度測った値で、標準偏差は求めていません。方向は確かでも、小数点以下は信じないほうがよいです。
現場での姿
今回の測定は、LabHubの本番のllmデプロイを変更していません。適用するには、イメージとサーバー引数と環境変数を変える必要があり、それは別の作業です。検討するときにレポートが指摘した条件は、文書の合計トークンがサーバーログのGPU KV cache sizeより大きくなければ差が見えない、ということです。同じ文書が再度読まれなければ、読み戻すものがないという点も同じです。CPU段階の容量は、トークンあたりのKVサイズから計算して決める必要があります。この測定の1.5GBは、文書10個分(約1.3GB)を収められたため通用しました。
次のラボですること
GPUなしで、純粋なPythonでこの流れを自分で確認します。LRUキャッシュを作り、容量21,520トークンに文書10個を同じ順序で2回読み込ませて、ヒットが0であることを確かめます。逆順に読むとどの文書がヒットするかを見て、容量をどこまで増やせばヒットが生じるかを探します。次に、GPUとCPUの2段キャッシュを作って、2回目のパスがすべてヒットすることを確認し、CPUの容量が足りないときを予測します。最後に、上の表の元データであるリクエスト別の測定ファイルを読み、p50と平均と倍率を自分で計算します。