推論基礎設施
SageMaker、プロンプトプレフィックスに基づくルーティングでKV cacheを維持――Llama 3.1 70Bの最初のトークンまでのレイテンシを最大77%削減
SageMakerリアルタイム推論エンドポイントに`PREFIX_AWARE`が追加され、同じシステムプロンプト、ドキュメント、会話履歴を持つリクエストが可能な限り同一インスタンスに送られるようになった。AWS公式の7ノードテストでは、長いプレフィックスを持つワークロードで顕著な効果が示された一方、シリアライズ方法の違い、ホットプレフィックス、キャッシュ容量によって実際の効果は左右される。
Amazon SageMaker Inferenceにプレフィックス対応ルーティングが追加された。これは、複数インスタンスで運用するLLMサービスにおいて「フレームワークがprefix cachingをサポートしていても、ロードバランサーによってキャッシュが無効化される」という問題に対処するものだ。従来のランダムルーティングや未完了リクエスト数が最少のインスタンスを選ぶ方式では、同じシステムプロンプト、RAGドキュメント、会話履歴を持つリクエストが異なるGPUに分散される可能性がある。その結果、各インスタンスでprefillが再実行され、計算済みのattention key-value tensorを再利用できない。
新しい`PREFIX_AWARE`戦略は、リクエストの先頭部分に基づいて安定したマッピングを作成し、共通プレフィックスを持つリクエストを同じインスタンスへ送る。エンジニアは`PrefixLength`を設定でき、ネイティブのInvoke APIでは1,024～65,536 bytes、OpenAI互換インターフェースではメッセージ本文の文字数を基準に計算される。`ConcurrencyThreshold`は1～1,024に設定でき、対象インスタンスが過負荷になった場合は1回分のキャッシュヒットを犠牲にして、より空いているノードへリクエストを送る。マルチテナントサービスでは、`X-Amzn-SageMaker-Prefix-Aware-Id`または`prompt_cache_key`を使ってルーティンググループを分離できる。また、この機能はinference componentsおよび動的LoRA adapterとも併用可能だ。
AWSは7台の`ml.p5.48xlarge`、vLLM、Llama 3.1 70Bを使用してテストを実施した。8,000-tokenの共通プレフィックスを使用するシナリオでは、P50のTime to First Token（TTFT）が71～77%、P90が33～37%短縮された。KV cacheのヒット率は約25%から82%へ上昇し、スループットは15～16%向上した。一方、より短く長さも不均一なShareGPT型の会話では、スループットの向上は1.7～2.0%にとどまった。ルーティング自体による追加レイテンシは約1.3～1.9ミリ秒だった。
ただし、これはモデルの高速化を自動的に保証するものではない。コンテナ側でもprefix cachingを有効にする必要があるうえ、ネイティブAPIはリクエストのbytesを直接比較する。JSON内の空白、フィールド順序の違い、サンプリングパラメータを先頭に配置することなどにより、本来共有されるべきトラフィックが別々に扱われる可能性がある。プレフィックスを短く設定しすぎると、ホットスポットが生じることもある。デプロイ担当者は、実際のプロンプト分布に合わせて2つのパラメータを調整するとともに、キャッシュヒット率、オーバーフロー率、TTFTのテールレイテンシ、スケーリング後のキャッシュウォームアップ時間を継続的に監視すべきだ。