GKE Inference Gateway で予測レイテンシ ベースのルーティングを使用する

このドキュメントでは、GKE Inference Gateway 内で llm-d によって提供される予測レイテンシ ベースのルーティングを有効にして使用する方法について説明します。デフォルトでは、GKE Inference Gateway は、ロードシグナルとプレフィックス キャッシュ アフィニティ ヒューリスティックの組み合わせを使用してリクエストをルーティングします。予測レイテンシ ベースのルーティングでは、静的なヒューリスティックの重みが、ライブトラフィックで継続的にトレーニングされた XGBoost モデルに置き換えられるため、ワークロード パターンが変化しても、より正確なルーティングの決定が可能になります。

予測レイテンシ ベースのルーティングを使用する場合

この機能は、ワークロードに次の条件が適用される場合に最も効果的です。

• プロンプトと完了の長さのばらつきが大きい: リクエスト サイズが大幅に異なる場合、キューの深さだけではサーバーの負荷を適切に把握できません。レイテンシ予測ツールは、リクエストあたりの実際のプリフィルとデコードのコストを考慮します。
• リクエストごとのレイテンシ SLO: アプリケーションが 最初のトークンまでの時間（TTFT）または出力トークンあたりの時間（TPOT）のターゲットを 個々のリクエストで指定すると、スケジューラはルーティング時にこれらのターゲットを適用します。これは、候補となる各 Pod のヘッドルーム（予測レイテンシから SLO ターゲットを引いた値）を計算することで行われます。
• 静的な重み調整が難しい: トラフィック パターンが変化するたびに、キャッシュ アフィニティとロードシグナルのバランスを頻繁に再調整している場合、オンラインでトレーニングされたモデルが自動的に適応します。

予測レイテンシ ベースのルーティングの仕組み

このセクションでは、予測レイテンシ ベースのルーティングで使用されるアーキテクチャとスケジューリング パイプラインについて詳しく説明します。

アーキテクチャ

予測レイテンシ ベースのスケジューリングでは、EPP 自体とともに、2 つの追加のサイドカー コンテナが EPP Pod 内にデプロイされます。

llm-d EPP スケジューリング パイプライン

予測レイテンシ ベースのスケジューリングが有効になっている場合、EPP は次のシーケンスの構成可能なプラグインを使用して各リクエストを処理します。

1. predicted-latency-producer: 予測サーバーを呼び出して、InferencePool 内の候補となるすべての Pod の TTFT と TPOT の推定値を取得します。これは、各 Pod の現在の KV キャッシュ使用率、キューの深さ、プレフィックス キャッシュの一致スコア、受信リクエストの特徴に基づいて行われます。レスポンスがクライアントに返された後、プロデューサーは観測された TTFT とトークン間のレイテンシを新しいトレーニング サンプルとしてトレーニング サーバーに返します。

   • フォールバック動作: 予測サーバーに到達できない場合や、 予測サーバーがエラーを返した場合、EPP は KV キャッシュ使用率、キューの深さ、プレフィックス キャッシュの一致に基づく複合スコアに自動的にフォールバックします。

2. prefix-cache-affinity-filter: このフィルタは、Pod のプレフィックス キャッシュの一致スコアがアフィニティのしきい値（デフォルトは 0.80）を超える場合に、候補セットをキャッシュ ウォーム Pod に絞り込みます。このしきい値は、本番環境で観測された 2 つの母集団を分離します。1 つは、以前のターンから会話履歴がキャッシュに保存されている Pod、もう 1 つはそうでない Pod です。このフィルタは、イプシロン グリーディーなエクスプロイトと探索戦略を実装します。

   • エクスプロイト（デフォルト パス）: このパスはキャッシュ ウォーム Pod にルーティングされるため、 スコアリングではキャッシュの再利用に重点が置かれます。

   • 探索（確率が低い）: このパスは、構成可能な割合のリクエストでフィルタを完全にバイパスし、コールド Pod にキャッシュ エントリをシードして、キャッシュの断片化を防ぎます。

   • TTFT ロードゲート: エクスプロイト パスでも、最適なキャッシュ ウォーム Pod の予測 TTFT が最適な全体 Pod の TTFT を構成可能なしきい値（デフォルトは 5,000 ミリ秒）を超えている場合、アフィニティは破棄され、 候補セット全体が使用されます。

3. slo-headroom-tier-filter（SLO リクエストのみ）: リクエスト に SLO ヘッダーが含まれている場合、候補となる Pod をポジティブ ティア（SLO を満たすと予測される）とネガティブ ティア（SLO に違反すると予測される）に分割します。

4. latency-scorer: 候補となる Pod にスコアを付けます。SLO ヘッダーがない場合は、予測レイテンシが最も低い Pod が選択されます。SLO ヘッダーがある場合、スコアは headroomSelectionStrategy を使用してヘッドルーム（SLO から予測レイテンシを引いた値）に基づいて計算されます。

   • least（デフォルト）: ビンパッキング。ヘッドルームが最も小さいポジティブな Pod にルーティングし、使用率を最大化し、負荷の低い Pod を将来のトラフィック バーストに備えて解放します。
   • most: スプレッド。ヘッドルームが最も大きいポジティブな Pod にルーティングし、予期しない負荷の急増に備えて余裕を残します。

5. latency-slo-admitter （SLO リクエストのみ）: 候補となる Pod が SLO を満たすと予測されない場合、ターゲットに到達しないと予測されるリクエストで容量を消費するのではなく、破棄可能なリクエスト（優先度 0 未満）を拒否します。SLO ヘッダーがない場合、または SLO を満たす Pod が存在する場合、このフィルタは無効になります。

6. weighted-random-picker: スコアに対する重み付きランダム選択を使用して、最終的な Pod を選択します。これにより、負荷が分散され、スコアの高い Pod が優先されます。

ストリーミング モード

predicted-latency-producer プラグインは、streamingMode パラメータを使用して構成される 2 つのトレーニング モードをサポートしています。

• streamingMode: false （デフォルト）: エンドツーエンド（E2E）のリクエスト レイテンシでトレーニングします。 ワークロードでストリーミング レスポンスと非ストリーミング レスポンスが混在している場合や、リクエストごとの SLO の適用なしでレイテンシを認識するルーティングのみが必要な場合は、このモードを使用します。
• streamingMode: true: 個別の TTFT モデルと TPOT モデルをトレーニングします。TTFT は最初のストリーミング チャンクに記録されます。TPOT は後続のトークンでサンプリングされます。ワークロードが完全にストリーミングされ、意味のある x-slo-ttft-ms / x-slo-tpot-ms の適用が必要な場合は、このモードを使用します。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

• Google Kubernetes Engine API を有効にする。
Google Kubernetes Engine API の有効化
• このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。

• 必要に応じて、Compute Engine API、Network Services API、Model Armor API を有効にします。

  API へのアクセスを有効にするに移動し、手順に沿って操作します。

• GKE Inference Gateway のデプロイが機能していることを確認します。 GKE Inference Gateway をデプロイするをご覧ください。

• InferencePool が、同一の GPU タイプ、モデルの重み、サービング構成の均一な Pod セットを使用していることを確認します。

• GKE クラスタがバージョン 1.32.3 以降であることを確認します。

• Helm をインストールします。Helm インストール ガイドをご覧ください。

予測レイテンシ ベースのスケジューリングを有効にする

次の手順では、GKE Inference Gateway デプロイで予測レイテンシ ベースのスケジューリングを有効にする方法について説明します。

ステップ 1: 予測レイテンシを有効にして InferencePool をインストールまたはアップグレードする

latencyPredictor.enabled=true フラグは、トレーニング サーバーと予測サーバーのサイドカーを EPP Pod 内にデプロイし、スケジューリング プラグイン パイプライン全体を接続します。

helm upgrade --install INFERENCE_POOL_NAME \
  --set inferencePool.modelServers.matchLabels.app=MODEL_SERVER_LABEL \
  --set provider.name=gke \
  --set inferenceExtension.monitoring.gke.enabled=true \
  --set inferenceExtension.latencyPredictor.enabled=true \
  --version LLM_D_VERSION \
  oci://LLM_D_REGISTRY_PATH

次のように置き換えます。

• INFERENCE_POOL_NAME: InferencePool の名前（例: vllm-llama3-8b-instruct）。
• MODEL_SERVER_LABEL: モデルサーバー Pod の選択に使用されるラベルキー。
• LLM_D_VERSION: 使用する llm-d Helm チャートのバージョン。
• LLM_D_REGISTRY_PATH: llm-d OCI レジストリ パス。

ステップ 2: デプロイを確認する

EPP Pod がすべてのサイドカー コンテナが準備完了した状態で実行されていることを確認します。

kubectl get pods -l app=INFERENCE_POOL_NAME-epp

EPP Pod には、EPP 自体、トレーニング サーバー、1 つ以上の予測サーバーなど、すべてのコンテナが実行中または準備完了の状態であることが示されます。

ステップ 3: ベースライン リクエストを送信する

SLO ヘッダーを有効にする前に、標準の推論リクエストを送信して、ルーティングが機能していることを確認します。

curl -i -X POST GATEWAY_IP:PORT/v1/completions \
 -H 'Content-Type: application/json' \
 -H 'Authorization: Bearer $(gcloud auth print-access-token)' \
 -H 'x-prediction-based-scheduling: true' \
 -d '{
    "model": "MODEL_NAME",
    "prompt": "PROMPT_TEXT",
    "max_tokens": MAX_TOKENS,
    "temperature": "0"
 }'

次のように置き換えます。

• GATEWAY_IP: ゲートウェイ サービスの IP アドレス。
• PORT: ゲートウェイ サービスのポート番号。
• MODEL_NAME: 推論に使用するモデルの名前。
• PROMPT_TEXT: 入力プロンプト。
• MAX_TOKENS: 生成するトークンの最大数。

x-prediction-based-scheduling: true ヘッダーは、このリクエストを予測レイテンシ スケジューリング パイプラインにオプトインします。予測ツールのウォームアップ期間中、EPP はヒューリスティック ルーティングにフォールバックします。

ステップ 4: SLO を認識するリクエストを送信する（省略可）

リクエストごとの SLO の適用を有効にするには、TTFT と TPOT のレイテンシ目標ヘッダーを追加します。

curl -i -X POST GATEWAY_IP:PORT/v1/completions \
  -H 'Content-Type: application/json' \
  -H 'Authorization: Bearer $(gcloud auth print-access-token)' \
  -H 'x-prediction-based-scheduling: true' \
  -H 'x-slo-ttft-ms: 500' \
  -H 'x-slo-tpot-ms: 50' \
  -d '{
    "model": "MODEL_NAME",
    "prompt": "PROMPT_TEXT",
    "max_tokens": MAX_TOKENS,
    "temperature": "0",
    "stream": true
  }'

次のように置き換えます。

• GATEWAY_IP: ゲートウェイ サービスの IP アドレス。
• PORT: ゲートウェイ サービスのポート番号。
• MODEL_NAME: 推論に使用するモデルの名前。
• PROMPT_TEXT: 入力プロンプト。
• MAX_TOKENS: 生成するトークンの最大数。

リクエスト ヘッダー:

• x-prediction-based-scheduling: true: リクエストを予測レイテンシ スケジューリング パイプラインにオプトインします。
• x-slo-ttft-ms: 許容される最初のトークンまでの時間の最大値（ミリ秒単位）。
• x-slo-tpot-ms: 許容される出力トークンあたりの時間の最大値（ミリ秒単位）。

予測レイテンシ スケジューリングをモニタリングする

レイテンシ予測ツールが有効になっている場合、EPP は Cloud Monitoring を介して追加の指標を公開します。

予測サーバーをスケーリングする

EPP は、受信リクエストごとに候補となる Pod ごとに 1 回の予測呼び出しを行います。各予測サーバー インスタンスは、約 300 QPS の予測作業を維持します。

予測サーバー インスタンス数の目安:

latencyPredictor.predictionServerCount Helm 値を更新して、予測サーバー インスタンスを追加します。

制限事項

• 均一な InferencePool が必要: 単一プール内で GPU タイプ、モデル バリアント、サービング構成が混在している場合はサポートされません。
• ウォームアップ期間: XGBoost モデルでは、予測が正確になるまでに十分なリアルタイム交通情報サンプルが必要です。
• SLO の適用: 適用はルーティング レイヤでのみ行われます。モデルサーバーは、選択後に SLO ターゲットを超えるリクエストを終了しません。
• ステータス: この機能はプレビュー版です。十分なテストを行わずに、厳格な SLA 要件がある本番環境のワークロードに使用することをおすすめしません。

次のステップ

• GKE Inference Gateway の構成をカスタマイズする。
• GKE Inference Gateway のロールアウト オペレーションを実行する。