XLA Serving Engine logo
出典: opensource.googleblog.com

Googleは、Pythonで開発した機械学習モデルをC++で実行する「XLA Serving Engine（XSE）」を、Apache 2.0ライセンスでオープンソース化しました。Googleの測定では、短い推論リクエストの処理能力が、同社が最適化したPythonサーバーの2〜5倍となり、メモリー使用量は4分の1になりました。主な対象は、検索や推薦など、大量の短いリクエストを低遅延で処理する用途です。

検索や広告で毎秒600万件超を処理

推論とは、学習済みのモデルを使って入力に対する予測や判定を行う処理です。XSEは、機械学習の計算を各種プロセッサー向けにコンパイルするOpenXLAの技術群に属し、モデルの書き出し機能と推論用の実行環境を提供します。

Googleは2025年からXSEの中核部分を本番環境に導入しており、現在は毎秒600万件を超えるクエリーを処理しています。利用先にはYouTube Shorts、Google Ads、検索、Gemini Omniの一部が含まれ、社内で利用をさらに広げる計画です。

対象には、推薦や順位付け、画像処理、入力を特徴量に変換するエンコーダー型のモデルなどがあります。一方、長い文脈を扱い、出力を逐次生成する大規模言語モデルについて、Googleは専用のメモリー管理を備えるvLLMやSGLangの利用を推奨しています。

モデルを共通形式にまとめ、実行時にコンパイル

XSEでは、Pythonの数値計算・機械学習フレームワークであるJAXで作ったモデルを、「XLA SavedModel」という形式に書き出します。この形式にまとめて保存するのは、モデルの計算内容と、学習によって得た重みです。

計算内容の保存には「StableHLO」を使います。StableHLOは機械学習モデルの演算を表す仕様で、モデルを作るフレームワークと、実行用コードに変換するコンパイラーの間をつなぐものです。重みの保存には、学習状態などを保存する仕組みであるOrbaxのチェックポイントを利用します。

実行側では、C++で動くランタイムがXLA SavedModelを読み込み、コンパイラーのXLAを使って対象ハードウェア向けのコードに変換します。推論リクエストを処理する経路からPythonインタープリターを外す構成です。

対応先はCPU、NVIDIA GPU、GoogleのAI向けアクセラレーターであるCloud TPUです。単一ホストと複数ホストの双方に対応し、単独の推論エンジンとして使うほか、既存の配信システムへ組み込むこともできます。複数のリクエストをまとめる動的バッチ処理や、通信方式のgRPCを使った予測の提供にも対応します。

Pythonの負荷は計算の呼び出しや前後処理にもある

XSEが取り除くPythonの実行負荷には、モデルの数値計算を呼び出す処理も関係します。推論の処理時間に含まれるのは、数値計算そのものに加えて、Pythonのオブジェクト生成、フレームワークへの処理の振り分け、入力の前処理や出力の後処理、データ転送などです。

数値計算をGPUやネイティブコードで実行していても、CPU側から小さな計算を順に起動する間にGPUの待ち時間が生じる場合があります。C++への移行が主に減らすのは、こうしたPython側の処理負荷です。XSEはC++の実行環境からモデルを動かすことで、Pythonインタープリターと、Pythonの並列実行を制約するグローバルインタープリターロックの負荷を取り除く設計です。

Googleの短いリクエストを対象とした測定では、単純な構成のPythonサーバーに対して8〜36倍の処理能力を得たとしています。予測1回当たりのCPU時間も30〜70％減少したとしています。

バッチ処理と同時実行数で変わる処理能力

XSEの動的バッチ処理は、届いた複数のリクエストをまとめて実行する機能です。推論サーバーでは、バッチサイズと同時に送るリクエスト数の組み合わせによって、処理能力と応答時間が変わります。

NVIDIAのTriton Inference Serverを使った測定では、動的バッチ処理を使わない構成で、同時リクエスト数を1から2へ増やすと、毎秒の推論数は159回から204.8回へ増えました。同時リクエスト数を4へ増やした場合は毎秒199.6回となり、95％のリクエストが収まる応答時間は約20.5ミリ秒まで延びています。

同じモデルで動的バッチ処理を有効にし、同時リクエスト数を8にした測定では、毎秒426.6回、95％のリクエストが収まる応答時間は約19.5ミリ秒でした。XSEも、動的バッチ処理や起動時のウォームアップをモデルのパッケージ内で設定できる構成です。

TorchTPUとの統合でPyTorchへの対応を拡張

Googleは今後、PyTorchをTPUで動かすTorchTPUとの統合を計画しています。PyTorchの利用者がモデルをXLA SavedModelに書き出し、TPU上のC++実行環境で提供できるようにする方針です。

現在のXSEでは、PyTorchのモデルをtorchax経由でJAXへ変換して利用する経路も用意されています。また、データ処理とモデルの計算グラフをC++で組み合わせるための前処理機能も備えています。