[論文レビュー] Prompt Cache: Modular Attention Reuse for Low-Latency Inference

In Gim, Guojun Chen+4 more|arXiv (Cornell University)|Nov 7, 2023

Topic Modeling被引用数 12

ひとこと要約

Prompt Cache は再利用可能なプロンプトモジュールの注意状態を再利用することで、GPU/CPUでのTTFT速度を8×–60×向上させ、精度の低下なし。

ABSTRACT

We present Prompt Cache, an approach for accelerating inference for large language models (LLM) by reusing attention states across different LLM prompts. Many input prompts have overlapping text segments, such as system messages, prompt templates, and documents provided for context. Our key insight is that by precomputing and storing the attention states of these frequently occurring text segments on the inference server, we can efficiently reuse them when these segments appear in user prompts. Prompt Cache employs a schema to explicitly define such reusable text segments, called prompt modules. The schema ensures positional accuracy during attention state reuse and provides users with an interface to access cached states in their prompt. Using a prototype implementation, we evaluate Prompt Cache across several LLMs. We show that Prompt Cache significantly reduce latency in time-to-first-token, especially for longer prompts such as document-based question answering and recommendations. The improvements range from 8x for GPU-based inference to 60x for CPU-based inference, all while maintaining output accuracy and without the need for model parameter modifications.

研究の動機と目的

• システムメッセージ、テンプレート、文書など、繰り返されるプロンプトセグメントを活用して大規模言語モデルの推論待機時間を削減する動機付け。
• 再利用可能なプロンプトコンポーネントとその位置を明示的に定義する形式的なプロンプトモジュールスキーマ（PML）を導入する。
• 事前計算およびプロンプトモジュールの注意状態をキャッシュし、プロンプト間で再利用して注意計算を最小化する。
• モジュラーな注意再利用を維持しつつ出力品質を保ちつつ、複数モデル/データセットで顕著な待機時間削減を実証する。

提案手法

• プロンプトモジュールをプロンプトマークアップ言語（PML）とモジュラーな注意再利用を可能にするスキーマで定義する。
• CPUまたはGPUメモリに対してプロンプトモジュールの注意状態をエンコード、キャッシュする。
• 推論時に完全なプロンプトの注意状態を形成するため、キャッシュされたプロンプトモジュール状態と新たに計算されたセグメントを連結する。
• モジュール型の注意再利用に必要な断続的な位置IDをサポートするよう、トランスフォーマーアーキテクチャを適応させる。
• HuggingFace Transformers上に構築されたプロトタイプ実装を LongBench データセットを横断して Llama2、Falcon、MPT で評価する。

実験結果

リサーチクエスチョン

• RQ1 共通のテキストセグメントを共有する異なるプロンプト間で注意状態を再利用できるか？
• RQ2 正確性を損なうことなく注意再利用を最大化するために、再利用可能なプロンプトモジュールをどのようにエンコードして管理するか？
• RQ3 Prompt Cache が CPU および GPU 配備における待機時間（最初のトークンまでの時間）とエンドツーエンド生成に与える影響は？
• RQ4 プロンプトモジュールの事前計算された注意状態を保存する際のメモリオーバーヘッドはどれくらいか、モデルサイズとプロンプト長に伴ってどうスケールするか？
• RQ5 Prompt Cache は LongBench の多様なタスク（QA、要約、コードなど）で出力品質を維持するか？

主な発見

• Prompt Cache は最初のトークンまでの待機時間を大幅に短縮し、GPUで8×–10×、メモリ配置次第ではCPUで最大20×–70×の利得をもたらす。
• 出力精度は LongBench のタスクと複数モデルアーキテクチャ（Llama2、MPT、Falcon）に対してベースライン KV キャッシュと同等の水準を維持している。
• 長いプロンプトや大きなキャッシュセグメントがあるほど、自己注意が二次的コストになる一方、キャッシュメモリコピーコストは線形であるため、待機時間の恩恵は大きくなる。
• GPUメモリにプロンプトモジュールを格納すると高速化が高まるが、CPUメモリは非常に大きなモジュールライブラリに対してより大きな容量を提供する。
• PMLでプロンプトモジュールとユニオンをパラメータ化することで、再利用可能性を保ちながら柔軟で再利用可能・結合可能なプロンプト設計を実現できる。