SageMakerの新ルーティング機能でLLM初回応答を最大77%高速化 — 複数サーバーに分散してもKVキャッシュを無駄にしない仕組み

9月11日、AWSが「Reduce LLM latency with prefix-aware routing on Amazon SageMaker Inference」と題した記事を公開した。この記事では、Amazon SageMaker InferenceにおけるKVキャッシュを効果的に活用するプレフィックス対応ルーティング機能と、それによるLLM推論レイテンシの大幅削減について詳しく紹介されている。

マルチインスタンス環境でKVキャッシュが効かない問題

LLMを使ったアプリケーションでは、プロンプトの冒頭に固定のシステム指示やドキュメントを置き、末尾にユーザー入力を続けるパターンが一般的だ。カスタマーサポートボットであれば、「あなたはAnyCompanyのサポートエージェントです。以下がポリシーです…」という3,000トークンの固定部分に、ユーザーの50トークン程度の質問が続く。

vLLMやTensorRT-LLMといったLLMサービングフレームワークは、この問題に対してプレフィックスキャッシュ（prefix caching）を実装している。一度処理したプロンプトの冒頭部分のKV（Key-Value）ペアをキャッシュし、同じプレフィックスが来たら再計算をスキップする仕組みだ。

ところがインスタンスが複数台に増えると、この恩恵がほぼ消える。デフォルトのランダムルーティングでは、同じ3,000トークンのプレフィックスがインスタンスA、B、Cと分散されてしまい、どのインスタンスもキャッシュを温め続けることができない。

PREFIX_AWAREルーティングで何が変わるか

今回SageMaker Inferenceに追加されたプレフィックス対応ルーティング（PREFIX_AWARE）は、リクエストの冒頭部分を見て、同じ先頭内容を持つリクエストを常に同じインスタンスへ送る仕組みだ。これにより各インスタンスのKVキャッシュが実際に蓄積・再利用されるようになる。

組み込みの保護機構も2つある：

過負荷保護：特定プレフィックスのリクエストが集中してターゲットインスタンスが逼迫した場合、空きのあるインスタンスへオーバーフローさせる。キャッシュヒットは逃すが、1台への集中を防ぐ。
スケーリング時の安定性：インスタンスの追加・削除時も大部分のリクエストは従来と同じインスタンスへ流れ続ける。スケールアウトのたびにキャッシュが無効化されることを避ける設計だ。

ベンチマーク結果

Llama 3.1 70B Instructを7台のml.p5.48xlargeインスタンス上でvLLM（プレフィックスキャッシュ有効）を使って検証した結果は以下のとおり。

長文コンテキスト（8,000トークンの共有プレフィックス、1時間継続）：

P50 TTFT（初回トークン生成時間）：71〜77%削減
P90 TTFT：33〜37%削減
KVキャッシュヒット率：約25% → 82%
スループット：15〜16%向上

短文コンテキスト（ShareGPT形式の会話、30分）：

P50 TTFT：13〜16%削減
KVキャッシュヒット率：約30% → 80%
スループット：1.7〜2.0%向上

ルーティング処理自体のオーバーヘッドは1.3〜1.9ミリ秒。モデルのTTFTが63〜280ミリ秒の範囲であることを考えると無視できるレベルだ。また7台へのトラフィック分散は各インスタンス13.3〜15.4%と均等を保っており、ホットスポットは発生しなかった。

設定方法

エンドポイント設定時にルーティング戦略を指定するだけで、モデルコンテナやサービングフレームワーク側の変更は不要だ。

aws sagemaker create-endpoint-config \
  --endpoint-config-name example-llm-config \
  --production-variants '[{
      "VariantName": "AllTraffic",
      "ModelName": "example-llm-model",
      "InitialInstanceCount": 3,
      "InstanceType": "ml.p5.48xlarge",
      "RoutingConfig": {
          "RoutingStrategy": "PREFIX_AWARE",
          "PrefixAwareRoutingConfig": {
              "PrefixLength": 4096,
              "ConcurrencyThreshold": 10
          }
      }
  }]'

設定パラメータは2つ：

PrefixLength（1024〜65536）：ルーティング判定に使うリクエスト先頭部分の長さ。ネイティブInvoke APIではバイト数、OpenAI互換APIではメッセージテキストの文字数で計る。なお、APIによって計測単位が異なるため、実装時は元記事の該当APIのリファレンスで単位を確認することを推奨する。
ConcurrencyThreshold（1〜1024）：ターゲットインスタンスの同時処理リクエスト数の上限。超えた場合は別インスタンスへ流す。

呼び出し側のコードは変更不要。既存のInvokeEndpoint APIもOpenAI互換APIもそのまま使える。

使いどころ

元記事が挙げているユースケースは以下の4パターンだ：

RAGアプリケーション：複数ユーザーが同じドキュメントについて質問する場合、そのドキュメントが共有プレフィックスになる。
マルチターン会話：会話履歴が蓄積するほど共有プレフィックスが長くなり、キャッシュ効果が増す。
テンプレート型ボット・アシスタント：長い指示文を毎リクエスト先頭に付ける構成。
コード補完：ファイル内容をコンテキストとして含めるコーディングアシスタント。

運用上の注意点

サービングフレームワーク側でもプレフィックスキャッシュを有効化すること。vLLMは最近のバージョンでデフォルト有効になっているが、他フレームワークは明示的な設定が必要な場合がある。
ネイティブInvoke APIではリクエストのJSONシリアライゼーションを統一すること。空白やキー順序の違いでルーティングが分散してしまう。
PrefixLengthは短すぎると一部インスタンスに集中、長すぎるとtemperature等の差異でスコープが分散する。共有プレフィックスの長さに少し余裕を持たせた値から始めるのが無難だ。
機能を使うには最低2インスタンスが必要。1インスタンスの場合は設定しても意味がない。

なお、複数テナントが同じプレフィックスを共有する構成でキャッシュコンテキストを分離したい場合は、ネイティブAPIではX-Amzn-SageMaker-Prefix-Aware-Idヘッダー、OpenAI APIではprompt_cache_keyフィールドでテナントIDを渡すことで分離できる。

詳細はReduce LLM latency with prefix-aware routing on Amazon SageMaker Inferenceを参照していただきたい。