なぜ埋め込みサービングが再び難しくなったのか
埋め込みモデルは、テキスト・画像・音声といった非構造化データを dense vector に変換し、セマンティック検索・推薦・意図分類といった下流タスクを支える「翻訳器」です。プロトタイプ段階では小規模なテキスト埋め込みモデルを動かすだけで十分ですが、数百万クエリ規模になるとボトルネックの性質が一変します。
実務でよく遭遇する課題は次の3点です。
- アクセラレータ容量の弾力性：トラフィック変動に合わせてコンピュートを即座に増減できること
- コスト/性能効率：同じスループットをより安く実現すること
- 数値パリティ：異種ハードウェアバックエンド間で埋め込み結果が数学的に一致すること（特にエンタープライズ要件）
特に次世代の埋め込みモデルは、テキストで 4K+ トークン、マルチモーダルで 15K+ トークンという超長文コンテキストを処理する必要があり、従来のサービングスタックでは HBM が持ちません。本記事では、Google Cloud が vLLM にネイティブ TPU サポートを組み込む過程で解決した4つの最適化を、実務視点で解説します。（根拠資料：Google Developers Blog 原文）
最小構成：Qwen3-Embedding-8B を TPU 上で起動する
vLLM のネイティブ pooling runner を使えば、TPU 上でも dense pooling 出力を直接取得できます。以下は公式レシピからの最小実行例です。
from vllm import LLM
# Cloud TPU 上で vLLM のネイティブ pooling runner により Qwen3-Embedding-8B を初期化
llm = LLM(
model="Qwen/Qwen3-Embedding-8B",
runner="pooling", # dense pooling 出力を有効化
tensor_parallel_size=2, # TPU トポロジメッシュにシャーディング
max_model_len=16384,
max_num_batched_tokens=512,
dtype="bfloat16",
trust_remote_code=True
)
# 入力に対する dense ベクトル埋め込みを抽出
prompts = ["Enterprise-grade semantic retrieval on TPUs with vLLM."]
results = llm.embed(prompts)
embedding_vector = results[0].outputs.embedding
ポイントは runner="pooling" です。これを有効にしないと LLM の生成ヘッド側から出力が出てしまい、埋め込みベクトルは得られません。また tensor_parallel_size は TPU トポロジメッシュのサイズと必ず一致させる必要があります。
内部的に何を修正しているのか
TPU 上で vLLM を動作させるには、大きく4つのポイントに手を入れる必要がありました。
- 語彙行列パディング — TPU の MXU（Matrix Execution Unit）は、テンソル並列（TP）で語彙行列をメッシュにシャーディングする際に厳密な整除制約を課します。そのため All-Gather 実行中のテンソルアラインメントを保証するハードウェアセーフなパディング戦略を実装しています。
- 遅延ローダ互換性（Attribute Promotion） — vLLM は TPU 上でコールドスタート遅延とホストメモリピークを抑えるために lazy-loading を使用します。ここで unquantization パイプラインに attribute promotion を導入しないと、初期化が失敗します。
- JIT コンパイルキャッシュの事前ウォーミング — マルチプロセスデプロイにおける JAX/XLA コンパイルトラップを回避するため、推論前にコンパイルキャッシュをロックしておく必要があります。
- Chunked Prefill の状態保持（StepPool + CachedRequestState） — 超長文コンテキストでは pooling レイヤーで Chunked Prefill が必須ですが、ステップ境界をまたぐと pooling 状態が失われる問題がありました。ハイブリッド StepPool とメタデータのマイグレーションにより、preemption 後も状態が正しく累積されるようにしています。
注意事項：パリティ検証なしで本番投入しないこと
TPU に移行すると性能は向上しますが、数値パリティの検証は必須です。埋め込みは検索品質に直結するため、わずかな誤差でもランキングが変動します。
検証方法はシンプルです。TPU で生成したベクトル（vTpu）と XPU のゴールデンレファレンス（vRef）のコサイン類似度を計算します。
| 入力種別 | 合格閾値 |
|---|---|
| テキスト | ≥ 0.999 |
| マルチモーダル（テキスト＋画像） | ≥ 0.995 |
1.0 に近いほど、ハードウェアバックエンド間の数値パリティが確保できていることを意味します。マルチモーダルの閾値がやや低いのは、画像エンコーダ側の演算経路が長く、累積誤差の余地が大きいためです。
この技術の限界と注意点
- TPU 依存性：vLLM-TPU スタックに最適化が入っているため、他のアクセラレータへ移行する際は別途検証が必要です。
- JIT ウォーミングコスト：pre-warming を行わないと初回リクエストのレイテンシが跳ねます。オートスケーリング環境ではノード起動ごとにこのコストを支払います。
- パディングオーバーヘッド：語彙サイズが TPU メッシュサイズの倍数でない場合、パディングによるメモリ浪費が発生します。
- Chunked Prefill のチューニング：
max_num_batched_tokensを小さくしすぎると状態遷移のオーバーヘッドが増大します。
あわせて読みたい
- CSS 相対カラー（Relative Color）でアニメーションテーマリングを完結させる — フロントエンド側の最適化感覚も同時に養いたい方へ。
- Netflix の Cassandra ワイドパーティション戦略 — 大規模時系列データを扱うシステム設計の観点から参考になります。
日本の開発コミュニティにおける適用文脈
日本では Cloud TPU を直接プロビジョニングして運用しているチームはまだ多くありません。多くの場合、AWS/GCP の GPU インスタンスで埋め込みサービングを運用しています。そのため本記事の内容を「明日すぐ適用」するのは難しいかもしれません。ただし、以下のシナリオでは検討価値があります。
- RAG パイプラインで埋め込みコストが GPU 請求の大きな割合を占める場合 — TPU のコストパフォーマンスが有効に働く可能性があります。
- 国内 SI 案件でオンプレミス埋め込みサーバが要求される場合 — vLLM が既にデファクトであるため、移行パスが開けています。
- マルチモーダル検索を自社構築するコマース/メディア企業 — 15K トークンのコンテキストが必要になる瞬間が来ます。
次のステップ学習の方向性
- vLLM アーキテクチャの理解：pooling runner、scheduler、chunked prefill がどう噛み合っているかをまず把握しましょう。
- TPU プログラミングモデル：MXU、HBM、トポロジメッシュの概念を理解すると、上記の最適化がなぜ必要かが見えてきます。
- パリティテストの自動化：CI にコサイン類似度の回帰テストを組み込んでおけば、ハードウェア交換時に安全です。
- GKE Custom Compute Classes：オートスケーリングの優先度ルールで複数のアクセラレータ種別を跨ぐ戦略を設計してみてください。
結論として、超長文マルチモーダル埋め込みを本番運用するには「単にモデルを起動する」のではなく、ハードウェア制約 → フレームワークパッチ → 数値検証という3段階を必ず踏む必要があります。この流れを理解しておけば、TPU でも GPU でも、どのバックエンドへ移行しても応用が効きます。