ハイパースケールキャッシュから学んだ GPU 利用効率の高め方

超低レイテンシのシステムから得た教訓が、LLM 推論を変えています。高性能キャッシュと推論が交わるところで、GPU の利用効率を高める方法を探ります。

Khawaja Shams headshot Khawaja Shams

超低レイテンシのシステムから得た教訓が、LLM 推論のあり方を変えています。

普段は接点の少ない二つの分野、高性能キャッシュシステムと大規模言語モデルの推論が交わるところで、静かな変革が起きています。

Momento は、100 マイクロ秒未満で応答するよう設計された、世界最速のハイパースケールキャッシュを構築してきました。今、その低レイテンシ分散システムに関する数十年の経験を、AI インフラの喫緊の課題の一つに生かしています。GPU を稼働させ続け、コストの際限ない増加を防ぐことです。

キャッシュが持つ意外な力

GPU の話に入る前に、システムの観点でキャッシュが実際に何をするかを確認しましょう。優れたキャッシュは、データベースの利用効率を高め、ユーザーが答えを得るまでの時間を短縮し、必要なデータベースサーバーの総数を減らします。ユーザー体験にも、財務にも有益です。

しかしデータベースから AI 推論へ移ると、同じ原則がはるかに大きな意味を持ちます。推論の世界では、利用効率によって、そもそも製品の運用に必要な資金、GPU、電力を確保できるかが決まるからです。

KV キャッシュの問題

ここから技術的に興味深い話になります。大規模言語モデルは、プロンプトとその周辺のコンテキストを処理するとき、単にテキストを読むのではなく変換します。入力をデコードし、KV(key-value)キャッシュと呼ばれるテンソルにします。プレーンテキストでは 1 メガバイト未満かもしれないプロンプトが、計算済みのテンソルデータとしてギガバイト規模に膨れ上がります。この計算は実際に GPU サイクルを消費します。

複数ターンの会話を考えてみてください。質問して回答を受け取り、追加の質問をします。新しいリクエストが前のメッセージを処理したものとまったく同じ GPU に届き、しかもその GPU のローカルメモリに KV キャッシュが残っていなければ、GPU はすべてのテンソルを毎回ゼロから計算し直さなければなりません。目を離すたびにすべてのクエリを忘れるデータベースのようなものです。

計算済みの KV キャッシュがノードのローカル DRAM や NVMe にあれば、それを取り出すほうが GPU で再生成するよりはるかに低コストです。ここまでは順調ですが、この最適化は新たな問題を生みます。

負荷の偏りという落とし穴

KV キャッシュを使えるようにする単純な方法は、同じユーザーのすべてのリクエストを同じ GPU に送ることです。しかし、これでは負荷が大きく偏ります。一部の GPU に処理が集中し、他は遊んだままになります。初期の Web サーバーのアーキテクチャを悩ませたのと同じ問題で、症状も同じです。性能が低下し、最初のトークンまでの時間が延び、ユーザー体験が予測しにくくなります。

そこで、一部のリクエストを負荷の低い GPU に回して平準化したくなります。ただし落とし穴があります。推論フリートが、流入するリクエストに追いつくため既に高負荷で動いていると、新しい GPU にテンソルを一から再計算させることで、混雑による障害に陥る可能性があります。負荷の問題を解いた代わりに、計算の問題を作ってしまうのです。

グローバルな KV キャッシュは、この悪循環を断ち切るためのものです。

足りなかった要素、グローバル KV キャッシュ

発想は単純ですが、実装は容易ではありません。新しい GPU で KV キャッシュを再計算する代わりに、既にフリート内のどこかにある計算済みテンソルを、利用する価値があるほど速く取得できるようにします。

ここで、100 マイクロ秒未満のキャッシュを実現してきた Momento の経験が直接役立ちます。高性能な分散キャッシュを支える三つの原則こそ、推論フリートに必要なものです。

1. 適切なルーティング。 リアルタイムのシグナルを基に、各リクエストを最適な利用可能ノードへ送ります。負荷だけでなく、データの局所性も考慮します。必要な KV キャッシュは既にどこにあるでしょうか。

2. 適切なデータ配置。 必要になる可能性が最も高い場所にデータを置きます。そのためには、対象のワークロードの挙動や、システム内でのデータの流れを深く理解する必要があります。

3. 高速なデータ移動。 ノード間の KV キャッシュ移動に時間がかかりすぎれば、この仕組みは成り立ちません。転送は、ゼロから再計算するより速く、実行する価値がなければなりません。KV キャッシュごとにギガバイト規模となれば、本格的な技術課題です。

従来のキャッシュでは、これらの原則がデータベースのボトルネックを解消します。推論では、遊休状態の GPU をなくします。

Rust、ベンチマーク、初期の結果

この考えを厳密に検証するため、低レイテンシと高性能が求められるシステムに適した言語、Rust で新しいベンチマークツールを開発しました。初期実験の数値は有望です。GPU がより多くの時間稼働するようになり、最初のトークンまでの時間(TTFT)は 50% 超短縮しました。

業界全体にとっても大きな意味があります。これまで AI インフラへの投資は、より大きく速いチップや、高価で専用のネットワーク機器の導入に集中してきました。一方、フリート全体で KV キャッシュの配置、ルーティング、転送を調整するデータ移動層には、まだ十分な検討が及んでいません。今回の取り組みが対象にするのは、その不足です。

なぜ重要なのか

現在、GPU 時間は世界で最も高価な計算資源の一つです。利用効率のわずかな改善も、大規模になると財務と環境に大きな影響を与えます。幸い、この不足を埋めるための実証済みの専門知識は、従来の分散システム工学に既に存在します。

100 マイクロ秒のキャッシュの世界と、数秒単位の LLM 推論の世界は、表面上は大きく異なります。しかし、その根底では同じ問題を解いています。適切なデータを、適切な場所へ、適切な時点で、意味のある速さで届けることです。

私たちは数十年にわたり、キャッシュの配置、ルーティングロジック、サブミリ秒の転送速度に取り組んできました。その知見を、今度は推論に生かしていきます。

あわせて読みたい