SageMaker Inference、プレフィックス単位のルーティング戦略を追加

何が発表されたのか

AWSは、Amazon SageMaker Inferenceのリアルタイム推論エンドポイントに「prefix-aware routing」と呼ばれる新しいルーティング戦略を追加しました。リクエストの先頭部分を見て、同じ先頭を持つリクエストを同じインスタンスに継続的に送る仕組みです。狙いは、LLMのサービングフレームワークが持つプレフィックスキャッシュを実際に機能させることにあります。

モデルコンテナやサービングフレームワーク側の変更は不要で、処理はエンドポイントのルーティング層だけで完結するとされています。呼び出し方法も従来どおりで、InvokeEndpointとInvokeEndpointWithResponseStream、およびOpenAI互換のChat Completion APIをそのまま使えます。戦略はエンドポイント設定で指定し、更新によって切り替えられるため、モデルの再デプロイは要りません。

スケールすると崩れる仕組み

LLMに送るプロンプトは、指示文・参照文書・会話履歴といった固定部分と、ユーザー入力の可変部分に分かれることが少なくありません。原文の例では、カスタマーサポートbotのリクエストが3,000トークンの指示文と50トークンの質問で構成され、何百・何千というリクエストで同じ3,000トークンが毎回処理されていました。

vLLMやTensorRT-LLMには、同じ先頭部分のKVペアを再利用するプレフィックスキャッシュがあり、time-to-first-token(TTFT)を大きく減らせます。ところがインスタンスを複数並べると、既定のランダム分散によって同じプレフィックスが別々のインスタンスに散らばります。どのインスタンスも同じ先頭を十分な頻度で見ないためキャッシュが育たず、機能はあっても活かせない状態でした。

prefix-aware routingは、この「送り先の一貫性」をルーティング層で担保します。過負荷保護として、人気のあるプレフィックスの送り先が設定した同時実行数の上限に達している場合は、より空いているインスタンスへ回す仕組みもあります。インスタンスを増減したときも大半のリクエストは同じ送り先に届き続け、キャッシュが毎回無効化されにくい設計とされています。

ベンチマークの結果

検証はLlama 3.1 70B Instructをml.p5.48xlarge 7台に載せ、vLLMのプレフィックスキャッシュを有効にした構成で、既定のランダムルーティングと比較されました。単一モデルエンドポイント、推論コンポーネントエンドポイント、ネイティブのInvoke API、OpenAI互換APIをまたぐ16通りの構成を実行し、いずれも成功率100%だったと報告されています。

8,000トークンの共通プレフィックスを1時間継続する長い文脈の条件では、P50 TTFTが71〜77%減、P90 TTFTが33〜37%減、スループットが15〜16%増となりました。KVキャッシュのヒット率は約25%から82%へ上昇しています。共通部分が長いほどスキップできる計算が増えるため、効果が大きくなると説明されています。

文脈が短い場合は効果が小さくなります。可変長のShareGPT形式の会話を30分回した条件では、P50 TTFTが13〜16%減、P90 TTFTが24〜37%減、スループットは1.7〜2.0%増、ヒット率は約30〜80%でした。ルーティング処理の追加コストは1リクエストあたり1.3〜1.9ミリ秒で、テスト時のモデルTTFTが63〜280ミリ秒だったことからすると小さいとされています。7台への配分も各13.3〜15.4%と偏りは1%以内に収まりました。

3つのルーティング戦略

今回の追加により、SageMaker Inferenceのリアルタイムエンドポイントが選べるルーティング戦略は3つになりました。既定のRANDOMはリクエストを全インスタンスへ均等に配り、LEAST_OUTSTANDING_REQUESTSは実行中のリクエストが最も少ないインスタンスへ送ります。

新しいPREFIX_AWAREは、同じプロンプト先頭部分を持つリクエストを同じインスタンスへ送ります。戦略は本番バリアント単位で設定でき、エンドポイント設定を更新するだけでモデルを再デプロイせずに切り替えられるとされています。

効果が出るのは、リクエストの先頭が共通しているワークロードです。原文では、同じ文書を前置きするRAG、履歴を積み上げるマルチターン会話、定型の指示文を毎回送るテンプレート型bot、同じファイル内容を文脈に含めるコード補完が例として挙げられています。

導入時の注意点

前提として、コンテナ側でプレフィックスキャッシュを有効にしておく必要があります。vLLMでは最近のバージョンで既定で有効ですが、他のフレームワークでは明示的な設定が必要になる場合があります。またインスタンスが1台ではどの戦略でも送り先は同じになるため、効果はありません。

挙動はPrefixLength(1024〜65536)とConcurrencyThreshold(1〜1024)の2つのパラメータで制御します。PrefixLengthはネイティブInvoke APIではリクエスト本文の先頭からのバイト数、OpenAI互換APIでは抽出されたメッセージ本文の文字数として扱われます。AWSは、共通プレフィックスの長さに少し余裕を足した値から始めることや、JSONの空白やキー順などのシリアライズを統一することを推奨しています。

マルチテナントでキャッシュの文脈を分けたい場合は、ネイティブAPIではX-Amzn-SageMaker-Prefix-Aware-Idヘッダ(最大64 ASCII文字)、OpenAI互換APIではprompt_cache_keyフィールドを指定します。推論コンポーネントエンドポイントや動的LoRAアダプタにも対応し、LoRAではアダプタが読み込まれたインスタンス群の中でプレフィックスに基づく選択が行われると説明されています。