分離型推論、第2回:デコードを止めずに KV キャッシュを移動する

LLM 推論の大きな改善は、GPU の高速化ではなく、KV キャッシュの移動を制御し、計算を繰り返さず再利用することから生まれます。

Hien Luu Hien Luu
個別のプリフィルとデコード処理を、HBM、DRAM、SSD にまたがる共有 KV ストアと比較する図

第1回では、分離構成の手間が見合う条件と、リクエストが適切な GPU を見つける方法を扱いました。ルーティングが終わると、次の問題は、デコードの流れを止めずに KV キャッシュをプリフィルからデコードへ移動することです。

高速な転送方式は、バイト列をプリフィルからデコードへ運びます。しかし、いつ送信を始めるか、途中でキャッシュをどこに置くか、リクエスト間でどう再利用するかは教えてくれません。これらはオーケストレーションの問題であり、実用上の性能の大半を左右します。

最初の手法は、層ごとのストリーミングです。Mooncake はこれを「Layer-wise Prefill」と呼びます。層0の KV キャッシュの計算が終わると、プリフィル GPU がまだ層 N を計算している間に、その層のブロックをデコードノードへ送り始めます。転送を計算の背後に隠し、そうしなければ生じる停止を避けることで、長いプロンプトの見かけの転送レイテンシを大幅に下げます。

2つ目は、KV キャッシュを階層型ストアとして扱うことです。ホットなブロックは HBM、ウォームなブロックは CPU DRAM、コールドなブロックはネットワーク接続の SSD に置きます。Mooncake Store は内容のハッシュをキーにし、LRU で追い出し、ホットなブロックを複数ノードへ複製します。人気のあるシステムプロンプトが1か所に集中してボトルネックになるのを防ぎます。DeepSeek の Fire-Flyer File System(3FS) はさらに進め、RDMA 越しの NVMe SSD によりノード内 DRAM への依存をなくします。180ノードの3FSクラスターは、KV 参照で合計6.6 TiB/s、クライアントあたり40 GiB/sの読み取りスループットを実現します。「キャッシュがディスクにある」ことが、予備手段ではなく実用的な選択になる速度です。

3つ目は、受け渡し方法そのものを選ぶことです。DistServe はプル方式で、デコードが必要に応じて取得し、プリフィルのメモリを待ち行列のバッファーとして使います。Mooncake の Conductor はプッシュ方式で、プリフィルが各層をデコードへ送り、メモリを解放します。3FS ベースの設計は共有ストレージを使います。プッシュはプリフィルプールのメモリ圧力と引き換えに、デコード側のレイテンシを最小化します。プルは、仕事が待機する場所にメモリ圧力をとどめます。共有ストレージは、ネットワーク1ホップと引き換えに両者を切り離します。

Perplexity の KV Messenger は、本番での実現例です。libfabric 経由の RDMA 上に構築され、各層の出力射影後に増えるカウンターをポーリングします。出力射影はテンソル並列のランク間でリダクションを行うため、カウンターが暗黙に同期を表し、CUDA グラフを崩さずに層ごとの完了を追えます。カウンターが進んだ瞬間に RDMA 書き込みが始まります。デコーダーには明示的な完了通知すら不要で、受信した RDMA 操作数を予定総数と照合します。DeepSeek-R1 の提供では、プリフィルとデコードを混在させると割り込みのため50 TPSを超えにくかったのに対し、分離後は TTFT の100ms増加と引き換えに、1つのプリフィラーが3つのデコーダーを90+ TPSで飽和させ続けました。

その見返りが再利用です。一般的なプロンプトの大半のトークンは、以前に誰かが処理しています。アプリの全ユーザーが共有するシステムプロンプト、継続中の会話の過去のターン、few-shot テンプレートの定型文などです。KV ストアが内容のハッシュで管理され、HBM、DRAM、SSD に階層化されると、それらを再びプリフィルする必要がなくなります。各ターンでやり直していた計算コストを、同じプレフィックスを持つすべてのリクエストで分担できます。

これを層ごとのストリーミングと組み合わせると、ワークロード全体の形が変わります。Moonshot AI が Kimi で1日1,000億トークン以上を処理し、同じ A800 ハードウェア上の非分離構成より115%多いリクエストを扱うのは、この仕組みによります。オーケストレーションが、計算の問題だったものをキャッシュの問題へ変えるのです。

次回

層ごとのストリーミング、階層型ストレージ、工夫された受け渡し方式は、すべて、その下のデータプレーンが数ギガバイトをミリ秒単位で GPU 間転送できることを前提にしています。そのための適切な転送方式を選び、当然に思える選択肢が不十分な理由を理解すること自体が、別の課題です。

次の第3回: ネットワーク基盤の準備が整っていないかもしれない理由。

あわせて読みたい