Managed Lustre の外部 KV キャッシュで AI 推論を迅速化

※この投稿は米国時間 2025 年 11 月 1 日に、Google Cloud blog に投稿されたものの抄訳です。

AI 推論インフラストラクチャの需要は増加の一途をたどっており、市場における AI 推論インフラストラクチャへの支出額は、間もなくモデル自体のトレーニングへの投資額を上回ると予想されています。この成長は、より豊かなエクスペリエンスに対する需要、特により大きなコンテキスト ウィンドウのサポートとエージェント AI の台頭によって促進されています。組織がコストを最適化しながらユーザー エクスペリエンスの向上を目指すには、推論リソースを効率的に管理することが最重要です。

大規模モデルの推論に関する実験研究によると、Google Cloud Managed Lustre などの高性能ストレージ上の外部 Key-Value キャッシュ（KV キャッシュまたは「アテンション キャッシュ」）を使用すると、総所有コスト（TCO）を最大 35% 削減できます。これにより、組織はプリフィル コンピューティングを I/O にオフロードすることで、43% 少ない GPU で同じワークロードを処理できます。このブログでは、長いコンテキストの AI 推論の管理における主な課題について説明するとともに、このような大幅なコスト削減と効率化を実現するために必要な高性能外部ストレージ ソリューションを Google Cloud Managed Lustre がどのように提供するについて、詳しく説明します。

KV キャッシュについて

推論フェーズにおいて、KV キャッシュは、Transformer ベースの大規模言語モデル（LLM）を効率的に運用するための重要な最適化手法です。

Transformer の重要なイノベーションは、シーケンシャル処理（繰り返し処理）を完全に排除したことであり、これはセルフアテンション メカニズムを導入することで実現されました。このメカニズムにより、シーケンス内のすべての要素が、自要素を他のすべての要素と瞬時に動的に比較し、関連性を評価できるようになりました（グローバルな一括評価）。このセルフアテンション メカニズムでは、モデルはシーケンス内の先行するすべてのトークンのキー（K）ベクトルと値（V）ベクトルを計算します。推論フェーズで次のトークンを生成するために、モデルには先行するすべてのトークンの K ベクトルと V ベクトルが必要であるためです。

ここで KV キャッシュが役立ちます。KV キャッシュは、（「プリフィル」ステージと呼ばれる）最初のコンテキスト処理後にこれらの K ベクトルと V ベクトルを保存します。これにより、後続のトークンを生成する際に、冗長でコストのかかるコンテキスト シーケンスの再計算を回避できます。この再計算を排除することで、KV キャッシュは推論プロセス全体を大幅に高速化します。小規模なキャッシュであれば、高帯域幅メモリ（HBM）またはホスト DRAM に収まる可能性があります。単一のマルチ アクセラレータ サーバーで数 TB までのメモリを利用できる場合もありますが、メモリ容量を超える複数の同時ユーザーのコンテキストの KV キャッシュを管理するには、外部ストレージ ソリューションまたは階層型ストレージ ソリューションが必要になることがよくあります。

ただし、このような大規模なコンテキストにより、AI モデルが大きなコンテキスト ウィンドウを処理する際に実行する「プリフィル」計算が非常に高価になる可能性があります。

• また、10 万トークン以上の大規模なコンテキストの場合、プリフィル計算により、最初のトークンまでの時間（TTFT）が数十秒に増加する可能性があります。

• プリフィル計算には、多数の浮動小数点演算（FLOP）が必要です。KV キャッシュを再利用することで、これらのコストを節約し、アクセラレータで追加のリソースを利用できるようになります。

エージェント AI の成長により、長いコンテキストを管理するという課題はさらに深刻化する可能性があります。単純な chatbot とは異なり、エージェント AI はなんらかのアクションを取ることを目的として構築されています。会話を行うだけでなく、問題を積極的に解決し、ユーザーに代わってタスクを完了します。これを行うために、幅広いデジタルソースからコンテキストを積極的に収集します。たとえば、エージェント AI は、フライトのリアルタイムのデータを確認したり、データベースから顧客の履歴を取得したり、ウェブでトピックを調査したり、独自のファイルに整理されたメモを保存したりする場合があります。そのため、エージェント AI は環境を深く理解できますが、多くの場合、コンテキストの長さとそれに関連する KV キャッシュのサイズが増加します。

パフォーマンス コストを大規模に管理する鍵は、アクセラレータを可能な限り最大限に活用することにあります。高性能なスケールアウト ストレージでは、アクセラレータあたりのスループットを向上させるため、リソース要件が軽減されます。

Google Cloud Managed Lustre 上の外部 KV キャッシュ

Google は、外部 KV キャッシュの主要なストレージ ソリューションとして Google Cloud Managed Lustre を推奨しています。GPU に関しては、Lustre はローカルに接続された SSD を利用します。ローカル SSD を利用できない TPU に関しては、Lustre の役割はさらに重要になります。

Google の Danna Wang による最近の LMCache ブログ投稿「LMCache on Google Kubernetes Engine: Boosting LLM Inference Performance with KV Cache on Tiered Storage」では、ホストレベルでのオフロードの基本的な価値が示されています。Google の Managed Lustre 戦略は、このホスト オフロードのコンセプトを次のレベルにまで進化させたものです。ローカル SSD と CPU RAM はノードローカルの階層として効果的ですが、サイズが固定されており、共有できません。Managed Lustre は、大規模な高スループット外部ストレージとして機能する並列ファイル システムを提供します。そのため、キャッシュがホストマシンの容量を超える、大規模なマルチノードかつマルチテナントの AI 推論ワークロードに最適なソリューションです。

Managed Lustre のパフォーマンス向上によって TCO を削減できる例を以下に示します。

• 50,000 トークンのコンテキストと高いキャッシュヒット率（約 75%）での実験では、Managed Lustre を使用すると、ホストメモリのみで KV キャッシュを使用した場合と比較して、合計推論スループットが 75% 向上し、最初のトークンまでの平均時間が 44% 短縮されました（詳細は以下を参照）。

• TCO 分析の結果、外部ストレージを利用しないワークロードと比較して、A3-Ultra VM と Managed Lustre を活用し、1 秒あたり 100 万トークン（TPS）を処理するワークロードに外部アテンション/KV キャッシュを使用すると、35% のコスト削減が実現することがわかりました。

• Google の実験では、構成のチューニングと、より多くの I/O 並列処理を実行できるようにする KV キャッシュ ソフトウェアの改善により、Managed Lustre で推論パフォーマンスを大幅に向上させることができることが実証されました。

総所有コスト: 分析

KV キャッシュ ソリューションを評価する際は、コンピューティングとストレージの費用だけでなく、運用にかかる費用と見込まれる節約額を含む TCO を考慮することが重要です。Google の分析によると、Managed Lustre 上に構築されたような、高パフォーマンス ストレージを基盤とする KV キャッシュは、純粋なメモリベースのソリューションと比較して、TCO の面で大きなメリットがあります。

コストの節約

増分ストレージ費用を考慮すると、100 万 TPS を処理するファイル システムを基盤とする KV キャッシュ ソリューションの TCO は、メモリのみのソリューションと比較して 35% 低くなると予測されます。そのため、大規模な AI 推論の導入において、よりスケーラブルかつ経済的に実行可能なオプションとなります。

TCO の主なメリットは、高価なコンピューティング リソースをより効率的に利用できることにより実現されます。KV キャッシュを高性能ストレージ ソリューションにオフロードすることで、アクセラレータあたりの推論スループットを向上させることができます。つまり、同じワークロードに必要なアクセラレータの数が減ります。特定の秒間クエリ数を 43% 少ないアクセラレータで処理できるため、直接的なコスト削減につながります。

TCO モデルの前提条件

TCO の計算には、いくつかの主要なコンポーネントが含まれています。

• ストレージ費用（正規価格）: これは、Managed Lustre の費用です。テストでは、TiB あたり 1,000 MB/秒のパフォーマンス階層を使用しました。TCO モデルには、100 万 TPS の目標レートを達成するのに十分な Lustre 容量（73 台の A3-Ultra マシン、マシンあたり 18 TiB の Lustre 容量）が含まれます。

• コンピューティング費用（正規料金）: A3-Ultra VM はそれぞれ 8 個の H200 GPU と 8 個の 141 GB HBM を搭載します（スポット料金はさらに低くなります）。

パフォーマンスのベンチマーク

Google Cloud Managed Lustre は、最先端の LLM に必要な高性能 I/O を提供できることが実験で実証されました。これらの実験では、Google Cloud A3-Ultra マシン（8x H200、8x 141GB HBM）で Deepseek-R1 を使用しました。また、50,000 トークンのコンテキストと高いキャッシュ ヒット率（約 75%）で、合計 KV キャッシュサイズが約 3.4 TiB の合成サービング ワークロードを実行し、メモリのみのベースラインでは、KV キャッシュに 1 TiB のホストメモリを使用しました。さらに、I/O 並列処理を多くした場合と少なくした場合の 2 つのバリエーションで Managed Lustre をテストしました。I/O の並列処理をより多く行うため、32 個の I/O ワーカー スレッドを利用して、Lustre から KV キャッシュデータを並列で読み取りました。

• Lustre を使用すると、ホストメモリのみで KV キャッシュを使用した場合と比較して、合計推論スループットが 75% 向上し、最初のトークンまでの平均時間が 44% 短縮されました。

推論ワークロードを最適化するには

長いコンテキスト ウィンドウの容量制限の問題を解決し、大規模な LLM で大幅なパフォーマンス向上を実現する外部 KV キャッシュ ソリューションを使い始めるには、次の手順を実行します。

1. インフラストラクチャをプロビジョニングし、Managed Lustre インスタンスを作成する

• 最適な低レイテンシ アクセスを実現するため、Lustre ファイル システムを利用するアクセラレータ（GPU または TPU）と同じリージョンとゾーンにプロビジョニングします。

• 推論エンジンをデプロイする: 外部 KV キャッシュまたは Paged Attention アーキテクチャをサポートする vLLM などの高パフォーマンスの推論サーバーまたは同様のフレームワークを使用して、LLM をデプロイします。

2. パフォーマンスを重視して構成する

Managed Lustre をマウントしたら、高性能ストレージを活用するように推論エンジン ソフトウェアを構成する必要があります。

• ダイレクト I/O を実装する: o_direct フラグを使用して Managed Lustre にアクセスするようにアプリケーションを構成します。これにより、汎用ファイル システム キャッシュがバイパスされ、推論エンジンが重要なホストメモリをより効果的に管理できるようになります。

• I/O 並列処理を調整する: 推論 KV キャッシュ ソフトウェアによっては、デフォルトのストレージ I/O 並列処理の設定が最適ではない場合があります。パフォーマンスを最大化するには、KV キャッシュ ソフトウェアをチューニングして、並列処理を強化した状態で KV チャンクファイルを読み取る必要がある場合があります。

次のステップに進むには、Managed Lustre の使用を開始する方法に関するドキュメントをご覧ください。

-上級ソフトウェア エンジニア、Kai Shen

-シニア プロダクト マネージャー、Barak Epstein