技術要約: PTStore (Prefix Tensor Store)

問題提起

大規模言語モデル（LLM）の推論ワークロードは、エネルギー消費とリソース需要の両面において、トレーニングを上回るハイパフォーマンス・コンピューティング（HPC）データセンターの主要な負荷となっています。LLMの推論は、プリフィル（prefill）（入力プロンプトを並列に処理するフェーズ）と、デコード（decode）（トークンを逐次的に生成するフェーズ）の2つのフェーズで構成されます。アテンション機構の冗長な計算を回避するため、システムはKey-Value (KV) キャッシュを利用して中間結果を保存します。

vLLMなどの最先端のランタイムは、単一のGPUまたはノード内でのKVキャッシュの最適化を行っていますが、スケールアップする際には以下のような重大な限界に直面します：

1. ノード間再利用の欠如: 既存のシステムは、分散コンピューティングノード間でのメモリ集約に失敗することがよくあります。あるノード上のリクエストが別のノード上のリクエストとプレフィックス（接頭辞）を共有している場合、第2のノードはキャッシュされたテンソルを再利用するのではなく、通常、プレフィックスを再計算してしまいます。
2. メタデータとレイテンシのボトルネック: 分散キャッシュを試みる手法（LMCache、EvoStoreなど）は、リモートメモリへのアクセスや複雑なメタデータ同期（例：単一ノードを超えてRadix-Attentionをスケーリングすること）による高いI/Oオーバーヘッドに苦しむことがよくあります。
3. メモリ制約: 個々のGPUメモリは、巨大なコンテキストウィンドウに対して不十分であり、ホストメモリやSSDへのオフロードは、キャッシュのメリットを打ち消すほどのレイテンシを導入します。

核心となる課題は、膨大なI/Oやメタデータのオーバーヘッドを引き起こすことなく、多くのコンピューティングノードに分散された多数のGPU間で、KVキャッシュ・プレフィックスの拡張可能かつ低レイテンシな再利用をどのように実現するかです。

手法: PTStore アーキテクチャ

PTStore（Prefix Tensor Store）は、KVキャッシュ・プレフィックスを分散・複製することでこれらの制限に対処するために設計された、分散型・複製型テンソルストアです。本システムはクライアント・サーバーモデルを採用しており、各コンピューティングノードでサーバーを実行し、ローカルのホストメモリとSSDを集約して、ローカルおよびリモートのGPUクライアントにサービスを提供します。

主要な設計原則

1. 増分テンソルストレージ（Trie構造）:

   • フルKVブロックを保存する代わりに、PTStoreは、新しいオブジェクトと、以前に保存されたオブジェクトの最長共通プレフィックス（LCP）との間の**増分差分（テンソル）**を保存します。
   • これにより、プレフィックスは、trie（トライ）のように、テンソルレベルの粒度で実装されながら、分岐する方向へ冗長性なく成長していくことが可能になります。
   • 統合されたメタデータ: 高コストな分散型trieの探索を避けるため、PTStoreはフラットなメタデータ構造を使用します。各オブジェクトのメタデータには、一意のテンソルIDのリストが含まれています。ロード操作は、これらのIDを反復処理して、レプリケーションキャッシュ内にローカルに存在するかどうかを確認します。存在しない場合は、その「オーナー」サーバーからそれらをリモートで取得します。

2. 分散階層型キャッシングとレプリケーション:

   • 所有キャッシュ (Owned Cache): 特定のサーバーが責任を持つ増分テンソルを保存します。
   • レプリケーションキャッシュ (Replication Cache): アクセス局所性を向上させるため、サーバー上に「ホット（人気のある）」なプレフィックスのコピーをローカルに保存します。
   • トレードオフ管理: システムは、所有キャッシュとレプリケーションキャッシュの間の構成可能な閾値を管理します。取得速度とストレージ容量のバランスをとるため、所有テンソル（低速なストレージへのフラッシュが必要）を退去させるよりも、レプリケーションされたテンソル（再取得が可能）を破棄することを優先します。

3. アクセスパターンを考慮したエビクション（退避）:

   • PTStoreは、LRU（Least Recently Used）ではなく、頻度ベースのエビクションポリシー（GDSFから適応）を利用します。これは、プレフィックス構造において、初期のテンソルほど頻繁にアクセスされるためです。
   • サイズと頻度のトレードオフを考慮し、小さな頻繁なテンソルが、大きな取得コストの高いテンソルを追い出さないようにします。

4. RDMAを意識した集約:

   • 分散を最小限に抑えるため、LCPに付加される増分は、オーナーサーバー上の単一の連続した領域に集約されます。
   • ロード操作は、データを転送する前に連続した領域にコピーするオーバーヘッドを避けるため、バルクRDMA（Remote Direct Memory Access）を使用して、分散されたセグメントを単一のRPC経由で並列にフェッチします。

主な貢献

1. 設計原則: 増分テンソルストレージ、統合メタデータ、およびプレフィックス・レプリケーションを統合した分散型リポジトリのための、一連の高レベルな設計原則。
2. PTStore プロトタイプ: これらの原則を実装した研究用プロトタイプ。vLLMのようなLLMランタイムとのシームレスな統合を可能にするC++低レベルAPIとPythonインターフェースを備えています。
3. 性能検証: 最先端のベースラインと比較して、I/Oオーバーヘッドとエンドツーエンドの実行時間を大幅に削減することを実証する広範な実験。

実験結果

著者らは、ALCF Polaris HPCテストベッド（560ノード、A100 GPU）上で、2つの抽出型QAワークロード（長コンテキストのWikiQA、および質問ボリュームが多いSQUAD）を用いてPTStoreを評価しました。使用されたLLMはMistral-7B-instruct-V2です。

ベースライン

• vLLM Vanilla: プレフィックス共有を行わない標準的なvLLM。
• vLLM Prefix: ノード内でのローカル・プレフィックス共有を行うvLLM。
• EvoStore: 増分ストレージとRDMAを使用しているが、ローカル・プレフィックス・レプリケーションを欠く分散型テンソルストア。
• PTStore: 分散型の認識とローカル・レプリケーションを備えた提案システム。

知見

• 弱いスケーラビリティ（8–32 GPUs）: PTStoreはEvoStoreおよびvLLM Prefixを大幅に上回りました。EvoStoreはリモート・プレフィックスのフェッチ時に高いRDMA I/Oオーバーヘッドに苦しんだのに対し、PTStoreのローカル・レプリケーションはこれを緩和し、TTFT（Time to First Token）において「切り離された優位性（detached advantage）」をもたらしました。
• シーケンス長のスケーラビリティ（1k–8k tokens）:
  • 短いシーケンス（1k）では、vLLMのローカルキャッシュが競争力がありました。
  • シーケンス長が増加するにつれ、PTStoreの優位性が増しました。8kトークンにおいて、PTStoreはvLLMのプレフィックス・キャッシングよりも約2倍速く、EvoStoreよりも20%高速でした。
  • コンテキストが長くなるほど、再計算またはリモートI/Oのコストがローカルのみのキャッシングのメリットを上回るため、性能差は拡大しました。
• 効率性の向上: 長いパッセージのQ&Aデータセットにおいて、PTStoreは、ノード間でメモリを集約せずKVキャッシュの再生を必要とするベースラインよりも、5〜6倍効率的に推論を実行しました。

重要性と主張

本論文は、PTStoreが、スケーラブルなLLM推論サービングにおける重要なギャップ、すなわち「分散ノード間でKVキャッシュ・プレフィックスを効率的に再利用できないこと」に対処していると主張しています。増分ストレージによる冗長性の最小化、迅速なクエリのための統合メタデータ、および局所性を最適化するレプリケーション戦略を組み合わせることで、PTStoreは以下を可能にします：

• クラスター全体のメモリを集約することによる、実効的なKVキャッシュサイズの数桁の拡張。
• 再計算のコストが高い長コンテキスト・ワークロードにおける、特に顕著なTTFTの削減。
• 通信のボトルネックやメタデータ同期の問題に悩まされる、従来の分散型アプローチを回避するスケーラビリティ。

著者らは、PTStoreをスケーラブルなAI推論への基礎的な一歩として位置付けており、今後の研究として、動的なメモリバランシング、機械学習を用いたエビクションポリシー、および実際の会話やコード補完のトレースを用いたLMCacheやMooncakeとのより広範なベンチマークへの注力を挙げています。