LabHub
はじめる
学ぶ 学習パス コース

LLMサービング

LRUと2段キャッシュのシミュレータでKVキャッシュのオフロードを確かめる

LabHubで続きを見る

目標

GPUのKVキャッシュの容量より多くの文書を交互に読み込ませたときに、LRUキャッシュがどのように崩れるかを自作のシミュレーターで確認し、CPU段階を加えると何が変わるかを見ます。最後に、実際の測定が残したリクエスト別のファイルから、倍率を自分で計算します。GPUは使わず、Pythonだけを使います。

なぜ重要なのか

長い文書を先頭に付けて質問だけを変えるリクエストでは、文書のKVキャッシュを再利用することが最大の節約になります。ところがGPUにはそのキャッシュを置く容量が小さいため、文書が容量より多くなると、古いものから追い出されます。この追い出しが循環アクセスに出会うと、容量が少し足りないだけでもヒットが0になる、というのがこのラボの1つ目の発見です。2つ目の発見は、追い出されるものをCPUメモリに退避しておけば、再計算する代わりに読み戻せるということで、2026-10-03の測定では、その差は最初のトークンまでの時間で576 ms対48.6 msでした。3つ目は限界です。シミュレーターは文書単位であり、測定は同時実行数1、モデル1.5B、GPU 1枚、合成文書でした。ラボのCPU容量56000トークンは、測定のCPU上限1.5GBを、文書あたりのKVサイズ(4,864トークンで0.1299GB)で割った近似値です。ステップ7(CPUの容量が足りないとき)はシミュレーターの予測であり、測定したものではありません。

ステップ

  1. /root/kvlab/kvsim.pyにLRUCache(capacity)を作ります。access(key, size)はヒットならTrueを返し、その項目を最も新しい位置に移します。ミスならFalseを返し、空きができるまで最も古いものから追い出したあとで保存します。容量より大きい項目は保存せず、ほかの項目も追い出しません。key in cacheは順序を変えず、cache.usedは保存されたサイズの合計(属性)、cache.keys()は古いものからのキーの一覧です。
  2. /root/kvlab/scenario.pyを作ります。引数は--gpu(デフォルト21520)、--docs(デフォルト10)、--tokens(デフォルト5011)、--passes(デフォルトはfwd,fwd,rev)です。文書番号0からdocs-1までがキーで、サイズはすべてtokensです。fwdパスは0から順に、revパスは逆順に読みます。標準出力にはJSONを1つだけ出力します: {"passes": [{"order": "fwd", "gpu_hits": 정수, "cpu_hits": 0, "misses": 정수, "hit_docs": [적중한 문서 번호를 읽은 순서대로]}, ...]}(プレースホルダーは整数と、ヒットした文書番号を読んだ順に並べた一覧で、オブジェクトはパスごとに1つです)。基本のシナリオを--passes fwd,fwdで実行した出力を、/root/kvlab/cyclic.jsonに保存します。
  3. --passes fwd,fwd,revで実行した出力を/root/kvlab/baseline.jsonに保存します。3つ目のパスのhit_docsにどの文書があるか、なぜそれらの文書なのかを読んで確認します。
  4. GPUの容量を21520、30000、40000、45000、50109、50110と変えながら--passes fwd,fwdで実行し、各容量での2回目のパスのヒット数(gpu_hitsとcpu_hitsの合計)を/root/kvlab/cliff.jsonに{"자리": 적중 수}の形で保存します(プレースホルダーは容量とヒット数です)。キーは文字列です。
  5. /root/kvlab/kvsim.pyにTwoTierCache(gpu_capacity, cpu_capacity)を追加します(LRUCacheはそのままにします)。属性gpuとcpuはそれぞれLRUCacheで、access(key, size)は"gpu"、"cpu"、"miss"のいずれかを返します。(1)GPUでヒットしたら"gpu"で、CPUには触れません。(2)GPUでミスしたがCPUにあれば"cpu"で、CPUのコピーを最新にし、GPUにも載せます。(3)どちらもミスなら"miss"で、GPUとCPUの両方に保存します。
  6. scenario.pyに--cpu(デフォルト0)を追加して、TwoTierCacheで実行します。出力のgpu_hitsはGPUのヒット、cpu_hitsはCPUから読み戻した数、hit_docsは両方を合わせた文書番号です。--cpuを指定しないか0にしたときは、ステップ2の結果と同じになる必要があります。--gpu 21520 --cpu 56000 --passes fwd,fwd,revの出力を/root/kvlab/offload.jsonに保存します。
  7. GPUの容量を21520にしたまま、CPUの容量を0、20000、30000、40000、50109、50110、56000と変えながら--passes fwd,fwdで実行します。各CPU容量での2回目のパスのミス数(misses)を、/root/kvlab/cpu_cliff.jsonに{"CPU 자리": 미적중 수}の形で保存します(プレースホルダーはCPU容量とミス数です)。キーは文字列です。
  8. 実測データをcp /opt/fixtures/kvoffload/ttft.jsonl /root/kvlab/ttft.jsonlでコピーします(内容は変更しません)。そして/root/kvlab/summarize.py <입력.jsonl> <출력.json>を作ります(プレースホルダーは入力ファイルと出力ファイルです)。入力は行ごとにJSONオブジェクト(config, pass, doc, doc_tokens, ttft_ms, total_ms)で、空行はスキップします。出力は{"baseline": {패스: {"n", "p50_ms", "mean_ms"}}, "lmcache": {패스: {...}}, "ratio": {패스: {"p50", "mean"}}}です(プレースホルダーはパス名です)。p50は中央値(個数が偶数なら中央の2つの値の平均)、msは小数第1位、ratioはbaselineをlmcacheで割った値で小数第2位です。パス名は入力にあるとおり(pass1-cold, pass2-replay, pass3-reverse)です。コピーしたファイルで実行し、結果を/root/kvlab/measured.jsonに保存します。

参考

LRUキャッシュを作る

/root/kvlab/kvsim.pyにLRUCache(capacity)を作成してください。access(key, size)はヒットならTrueを返し、その項目を最も新しい位置に移します。ミスならFalseを返し、空きができるまで最も古いものから追い出したあとで保存します。容量より大きい項目は保存せず、ほかの項目も追い出しません。key in cacheは順序を変えず、cache.usedは保存されたサイズの合計(属性)、cache.keys()は古いものからのキーの一覧です。

PythonのOrderedDictには、項目を末尾へ移すmove_to_endと、先頭を取り出すpopitem(last=False)があり、この作業に向いています。ヒットしたときに順序を変えるのを忘れると、LRUではなく入った順に追い出すFIFOになります。容量より大きい項目に出会ったときに先に空にしてしまうと、問題のない項目まで失います。

同じ順序で2回読む

/root/kvlab/scenario.pyを作成してください。引数は--gpu(デフォルト21520)、--docs(デフォルト10)、--tokens(デフォルト5011)、--passes(デフォルトはfwd,fwd,rev)です。文書番号0からdocs-1までがキーで、サイズはすべてtokensです。fwdパスは0から順に、revパスは逆順に読みます。標準出力にはJSONを1つだけ出力してください: {"passes": [{"order": "fwd", "gpu_hits": 정수, "cpu_hits": 0, "misses": 정수, "hit_docs": [적중한 문서 번호를 읽은 순서대로]}, ...]}(プレースホルダーは整数と、ヒットした文書番号を読んだ順に並べた一覧で、オブジェクトはパスごとに1つです)。基本のシナリオを--passes fwd,fwdで実行し、その出力を/root/kvlab/cyclic.jsonに保存してください。

パスごとに、同じキャッシュを1つ使い続けます(パスが変わってもキャッシュを空にしません)。ヒットした文書は、読んだ順にhit_docsに入れます。出力に説明文を混ぜると、採点ツールがJSONとして読めません。保存はpython3 scenario.py --passes fwd,fwd > cyclic.jsonのようにします。

逆順に読むと何が残っているか

同じシナリオを、最初、同じ順序で再度、逆順の順に読むよう--passes fwd,fwd,revで実行し、その出力を/root/kvlab/baseline.jsonに保存してください。3つ目のパスのhit_docsにどの文書があるか、なぜそれらの文書なのかを読んで確認してみてください。

新しいコードは必要ありません。ステップ2のscenario.pyを、引数だけ変えて実行します。3つ目のパスが始まるときにキャッシュに残っているものは何か、逆順に読むとそれをどの順序で見ることになるかを考えてみてください。

容量をどこまで増やせばヒットが生じるか

GPUの容量を21520、30000、40000、45000、50109、50110と変えながら--passes fwd,fwdで実行し、各容量での2回目のパスのヒット数(gpu_hitsとcpu_hitsの合計)を、/root/kvlab/cliff.jsonに{"자리": 적중 수}の形で保存してください(プレースホルダーは容量とヒット数です)。キーは文字列です。

容量を2倍近く増やしてもヒットが生じない区間があります。文書10個のサイズの合計がいくつになるかを計算し、それより1トークン足りない容量とちょうど合う容量を、実際に実行して比べてみてください。6回を手で実行せず、ループで回してJSONにしてください。

GPUとCPUの2段キャッシュ

/root/kvlab/kvsim.pyにTwoTierCache(gpu_capacity, cpu_capacity)を追加してください(LRUCacheはそのままにします)。属性gpuとcpuはそれぞれLRUCacheで、access(key, size)は"gpu"、"cpu"、"miss"のいずれかを返します。ルールは3つです。(1)GPUでヒットしたら"gpu"で、CPUには触れません。(2)GPUでミスしたがCPUにあれば"cpu"で、CPUのコピーを最新にし、GPUにも載せます。(3)どちらもミスなら"miss"で、GPUとCPUの両方に保存します。

段ごとにLRUCacheを1つずつ使えば足ります。ヒットしたかどうかを知るには、accessを呼ぶ前にinで先に確認する必要があります(accessはミスのときに保存までしてしまいます)。CPUの容量が0でも動作する必要があり、そのときは1段のキャッシュと同じになります。

CPUに退避すると2回目のパスはどうなるか

scenario.pyに--cpu(デフォルト0)を追加して、TwoTierCacheで実行してください。出力のgpu_hitsはGPUのヒット、cpu_hitsはCPUから読み戻した数、hit_docsは両方を合わせた文書番号です。--cpuを指定しないか0にしたときは、ステップ2の結果と同じになる必要があります。--gpu 21520 --cpu 56000 --passes fwd,fwd,revの出力を/root/kvlab/offload.jsonに保存してください。

runが作るキャッシュをLRUCacheからTwoTierCacheに変え、accessの戻り値3種類をそれぞれ数えれば足ります。前のステップのシナリオが壊れていないか、--cpuなしでも再度実行してみてください。得られた結果を測定と見比べてください。逆順に読むパスで、GPUヒットとCPUヒットがそれぞれ何個か、測定で約30 msと約48 msに分かれた文書数と合いますか。

CPUの容量が足りないとき

GPUの容量を21520にしたまま、CPUの容量を0、20000、30000、40000、50109、50110、56000と変えながら--passes fwd,fwdで実行してください。各CPU容量での2回目のパスのミス数(misses)を、/root/kvlab/cpu_cliff.jsonに{"CPU 자리": 미적중 수}の形で保存します(プレースホルダーはCPU容量とミス数です)。キーは文字列です。

ステップ4と同じループで、変える引数と読む値だけが違います。結果を見て、CPU段階に必要な容量が「GPUに入らなかった分」なのか「文書全体」なのかを判断してみてください。この値はシミュレーターの予測です。この測定では、CPU上限が文書をすべて収められる場合だけを測りました。

実測ファイルから倍率を計算する

実測データをcp /opt/fixtures/kvoffload/ttft.jsonl /root/kvlab/ttft.jsonlでコピーしてください(内容は変更しません)。そして/root/kvlab/summarize.py <입력.jsonl> <출력.json>を作成してください(プレースホルダーは入力ファイルと出力ファイルです)。入力は行ごとにJSONオブジェクト(config, pass, doc, doc_tokens, ttft_ms, total_ms)で、空行はスキップします。出力は{"baseline": {패스: {"n", "p50_ms", "mean_ms"}}, "lmcache": {패스: {...}}, "ratio": {패스: {"p50", "mean"}}}です(プレースホルダーはパス名です)。p50は中央値(個数が偶数なら中央の2つの値の平均)、msは小数第1位、ratioはbaselineをlmcacheで割った値で小数第2位です。パス名は入力にあるとおり(pass1-cold, pass2-replay, pass3-reverse)です。コピーしたファイルで実行し、結果を/root/kvlab/measured.jsonに保存してください。

configとpassの組ごとにttft_msを集めて、statistics.medianとstatistics.meanを使います。結果を読むときは、ratioが1より大きいか小さいかを見てください。最初のパスで1より小さい値は何を意味しますか。採点ツールは、文書の個数と順序が異なる入力でもこのスクリプトを実行します。