vLLM、Ironwood TPUで16K埋め込みをサポート、StepPoolが分割prefillの状態を保持

Google Cloudは、vLLM-TPUによる長いコンテキストの埋め込みサポートと、再現可能なデプロイ手順を公開した。対象モデルには、テキスト版のQwen3-Embedding-8Bとマルチモーダル版のQwen3-VL-Embedding-8Bが含まれる。今回の取り組みの焦点は、新たなモデル重みの追加ではない。自己回帰生成を中心に設計されてきたservingエンジンで、Ironwood TPU上のpooling、16K級の入力、マルチモーダルprefillを安定して実行できるようにすることだ。

最初の変更は、TPU MXUにおけるテンソルの割り切れ条件に対処する。モデルの語彙行列をtensor parallelで分割する際、vLLMはハードウェア要件を満たす安全なpaddingを追加し、形状の不整合によるAll-Gatherの失敗を防ぐ。2つ目の変更ではlazy loadingを強化し、サービスがリクエストを受け付ける前に、シャードを考慮したJAX／XLAコンパイルキャッシュをウォームアップする。これにより、マルチプロセス構成で最初のリクエスト群を処理するときに初めてコンパイルが発生するリスクを抑える。

さらに重要なのが、ハイブリッドStepPoolだ。長いシーケンスでは、prefill中にHBMを使い切らないよう入力を複数のchunkに分割する必要がある。しかし、埋め込みではシーケンス全体に対するpoolingが必要になる。途中の集約状態が単一のstepにしか紐付いていない場合、リクエストがpreemptionされたり再開されたりした際に状態が失われる可能性がある。新版では関連するmetadataを`CachedRequestState`へ移し、集約状態をchunk間で蓄積し、preemption後も引き継げるようにした。現時点でマルチモーダル版が分割するのはテキスト部分だけであり、画像パスは依然として容量上の制約を見極める必要がある。

公式の構成例では、単一の2×2×1 Ironwoodノード、4基のTPU、BF16、TP=4、および16,384-tokenのランダム入力を使用している。Googleのブログは83,996 token/s、5.13 request/sと報告する一方、GitHubの構成表では84,038.27 token/s、同じく5.13 request/sとなっている。このわずかな差は、異なる実行結果に基づく数値である可能性を示している。精度検証では、TPUと参照バックエンドが生成したベクトルのcosine similarityを使用し、しきい値をテキストで0.999、マルチモーダルで0.995としている。これはバックエンド間の数値的一致を確認できるものの、検索品質、テールレイテンシ、コストの面でGPUより優れていることを証明するものではない。エンジニアリングチームは次の段階として、自社の入力長分布、テキストと画像の比率、GKEのスケーリング挙動に基づいてテストを再実行すべきだ。