Perplexity、GPU 埋め込みスタック「Ivy」「Tulip」「ROSE」の詳細を公開

AI 検索製品の検索精度は、主に「埋め込みモデルの性能」と「インデックス全体での運用コスト」の二点によって制約されます。今週、Perplexity のエンジニアリングチームが公開した「Fast Embeddings on GPUs」では、後者、つまり pplx-embed や Perplexity Search、Computer、API プラットフォームで採用されているランキングモデルを支えるインフラストラクチャの詳細が明かされました。

Perplexity チームによれば、成熟した Hopper および Blackwell ハードウェア上での GPU における埋め込み推論は、エンジン間でほぼ標準化されています。その優位性は、モデル自体よりも、それを動かすランタイムとハッチ（基盤）にあります。具体的には CUDA グラフの管理、非同期結果追跡のための抽象化レイヤー、そして Rust で実装されたリクエスト処理パスです。

二つのトラフィックパターン、一つのエンジン

Perplexity は埋め込みサービスの提供を、主に二つのワークロードとして捉えています。一つはバッチ埋め込みで、ベクトルデータベースの構築や再インデックス時に発生します。ここではスループット最大化によるコスト削減が最優先されます。もう一つはオンライン埋め込みで、ユーザーからのクエリ処理時に発生します。ここでは短いクエリを高速に埋め込むことが求められます。そしてその中間に「スコアリング」があります。ベクトル検索の後に大量のドキュメントバッチをランキングする工程で、両者のバランスが取られています。

Perplexity が下した重要な決断は、専用の埋め込みエンジンを構築しなかったことです。埋め込みモデルは小型の Transformer であるため、バッチ埋め込みは計算リソースに制約される「prefill（事前填充）」に近く、オンライン埋め込み（数トークン程度）はメモリ帯域に制約される「decode（復号化）」に近い性質を持っています。そのため研究チームは、LLM スタンダードから流用した prefill と decode カーネルをそのまま再利用しています。

Ivy、Tulip、ROSE

一つのリクエストには、以下の三つのサービスが関与します：

Ivy は Rust で書かれた HTTP ゲートウェイです。JSON パース、トークン化、入力テンプレート処理、バッチ分割といった CPU 側の処理を担当し、リクエストをカスタム gRPC プロトコルに変換します。また、大規模なバッチリクエストをチャンクに分割してレプリカ間で負荷分散を行うため、生産環境でペイロードサイズがばらつくことによる負荷の偏りを解消しています。

Tulip は推論サーバーのインターフェースです。Rust、tokio、tonic を用いて構築された gRPC サーバーであり、リクエストのスケジューリングとバッチ処理を担当した上で、エンジンへ dispatch します。

ROSE（Runtime-Optimized Serving Engine）はモデル推論を実装するコンポーネントです。主に Python で書かれており、カーネルやレイヤー、モデル定義を提供します。CUDA グラフの管理を行い、Tulip に対して step() 関数を公開しています。

なぜスケジューラーはあえてシンプルなのか

Tulip はリクエストが蓄積される間、先着順でシーケンスを選択します。この単純なアプローチには明確な根拠があります。Perplexity が扱うような短いシーケンス長における小規模な埋め込みモデルでは、密結合層の線形コストがアテンションの二次コストを上回ります。つまり、レイテンシはシーケンス数ではなくトークン数にほぼ比例します。GPU を飽和させるバッチサイズ（例えばサブビリオンパラメータモデルで約 512 トークン）に達すると、さらに多くのシーケンスを詰め込んでも効率化にはつながりません。

CUDA グラフと LazyTensors

バッチサイズが小さい場合、CPU 側でのカーネル起動にかかる時間が GPU 実行時間を上回ってしまうことがあります。Perplexity ではすべての埋め込みモデルに対して全体モデルの CUDA グラフを構築し、すべての起動処理を単一のドライバ呼び出しに集約しています。埋め込みモデルは小規模であるため、GPU の処理コストが起動コストを上回る転換点は、数千トークン・数十シーケンスという大きなバッチサイズになって初めて訪れます。一部の注意機構（アテンション）実装では、動的なホスト側の入力に依存しているため全体グラフの構築を阻害しますが、Perplexity は FlashInfer への変更を上流に反映させることでこの問題を解決しました。

グラフは設定ごとに個別にキャプチャする必要があり、トークン数は 64 または 256 の倍数となるバケットにパディングされます。それでもモデルごとに数千のグラフを生成し、キャプチャには数分を要します。これを改善するため「遅延キャプチャ」を採用しました。各設定に対して最初のヒットで予備実行（ウォームアップ）を行い、2 回目のアクセス時にキャプチャと再生を開始する仕組みです。これにより起動時の p99 レイテンシは若干増大しますが、数分かかる予備処理を数時間にわたって分散させることが可能になります。

2 つ目の工夫が LazyTensor です。これはページロックされたホストバッファ、cudaMemcpyAsync 呼び出し、そして CUDA イベントを追跡する仕組みです。従来のように step() メソッドでデバイス側でブロックするのではなく、LazyTensor を返すことで、Rust の非同期タスクがバッチ N の完了を待っている間に CPU がバッチ N+1 の処理をキューイングできます。

カーネルの実装も依然として重要です

ROSE は、不規則な入力（ragged inputs）に対応する複数のアテンションバックエンドをサポートしています。具体的には FlashInfer 2、FlashInfer 3、そして FlashAttention 4 です。Perplexity チームによると、一般的には FlashAttention 4 の方が高速ですが、非常に長いシーケンス長において Qwen ベースのモデルでは FlashInfer 3 が上回るケースがあります。そのため、バックエンドの選択は状況に応じて個別に行われます。

特筆すべき点として、埋め込みモデル（embedding model）を推論する際、ROSE は KV キャッシュをインスタンス化しません。代わりにパディングを避けるために、不規則な入力に対応したアテンション変種に処理を委譲します。

ベンチマーク

Perplexity のベンチマークは、BF16 形式で実機上の重みと評価データから導かれた入力を用いて vLLM v0.22.0 と比較されています。ウォームアップ実行により、コサイン類似度の乖離が 0.1% 以内であることを確認しています。

以下の 4 つのベンチマークスイートが計測されました：

• 低遅延埋め込み（バッチサイズ 1；トークン数 128/512/4096）
• 低遅スコアリング（バッチサイズ 5/25/50、トークン数 512）
• 高スループット埋め込み（バッチサイズ 100、並行プロセス 4 つ）
• 高同時実行埋め込み（並行リクエスト数 1〜16、Ivy のトークナイゼーションやネットワークオーバーヘッドを含む）

主なポイント

Perplexity の埋め込みスタックは、専用のエンジンを別途動かすのではなく、LLM のプリフェッチ/デコードカーネルを流用しています。

遅延時間はシーケンス数ではなくトークン数に依存します。約 512 トークンで、10 億パラメータ未満のモデルが飽和状態になります。

モデル全体のカスタム CUDA グラフと遅延キャプチャにより、起動時のオーバーヘッドを削減しつつ、数分にも及ぶ起動時間を回避しています。

LazyTensor を採用することで、CPU でのバッチ準備処理と GPU 上の実行中ワークロードを並列化し、同期待ちによるブロックを防いでいます。

Ivy、Tulip、ROSE は内部システムです。外部からは Perplexity の Embeddings API を通じて pplx-embed にアクセスできます。

本記事「Perplexity Details Its GPU Embedding Stack: How Ivy, Tulip and ROSE Serve pplx-embed」は、MarkTechPost で公開されたものです。