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

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

予測レイテンシに基づくルーティングを使用する場合

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

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

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

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

アーキテクチャ

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

コンポーネント 説明
トレーニング サーバー EPP から受信した完了済みリクエスト サンプルで XGBoost TTFT モデルと TPOT モデルを継続的に再トレーニングします。スライディング ウィンドウで階層化されたバケットを使用するため、まれなトラフィック レジームが忘れられることはありません。更新されたモデルを共有ボリュームに書き込みます。
予測サーバー リクエストのホットパスで TTFT と TPOT の予測を EPP に提供します。共有ボリュームから最新のトレーニング済みモデルを読み取ります。水平スケーラビリティ - 各サーバー インスタンスは、約 300 QPS の予測作業を維持します。複数のインスタンスは、EPP の Go 統合プロキシによって負荷分散されます。このプロキシは、1 ミリ秒のウィンドウ内で同時予測リクエストをバッチ処理します。

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 つの母集団(以前のターンで会話履歴がすでにキャッシュに保存されている Pod と、そうでない 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 リクエストのみ): SLO を満たす候補 Pod が予測されない場合、ターゲットを満たさないと予測されるリクエストで容量を消費する代わりに、削除可能なリクエスト(優先度が 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: 許容可能な最大 Time-to-First-Token(ミリ秒単位)。
  • x-slo-tpot-ms: 許容可能な出力トークンあたりの最大時間(ミリ秒単位)。

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

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

指標 説明
inference_objective_request_ttft_seconds 実際の TTFT 分布(streamingMode=false の場合は E2E レイテンシ)。
inference_objective_request_predicted_ttft_seconds 予測された TTFT 分布(streamingMode=false の場合は E2E レイテンシ)。
inference_objective_request_tpot_seconds 実際の TPOT 分布。
inference_objective_request_predicted_tpot_seconds 予測された TPOT 分布。
inference_objective_request_ttft_slo_violation_total TTFT SLO 違反のカウンタ。

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

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

Prediction Server インスタンス数の概算の目安は次のとおりです。

クラスタ QPS(100 個の Pod) 予測サーバーが必要
最大 1,000 QPS 1 個のサーバー
最大 5,000 QPS 2 台のサーバー
最大 10,000 QPS 4 台のサーバー

latencyPredictor.predictionServerCount Helm 値を更新して、Prediction Server インスタンスを追加します。

制限事項

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

次のステップ