9月5日、MarkTechPostが「Perplexity Details Its GPU Embedding Stack: How Ivy, Tulip and ROSE Serve pplx-embed」と題した記事を公開した。この記事では、PerplexityがAI検索サービスの内部で動くGPUエンベディング推論スタック「Ivy」「Tulip」「ROSE」の設計と最適化手法について詳しく紹介されている。

本番AIサービスのMLOps内部構造がここまで公開されることは珍しい。Perplexityのエンジニアリングチームが公開したブログ記事「Fast Embeddings on GPUs」は、pplx-embedとランキングモデルをPerplexity Search・Computer・API Platformに提供するサービングインフラの設計を詳細に解説している。

3サービスで1リクエストを処理する

エンベディング推論パイプラインは3つのコンポーネントに分かれている。

• Ivy：RustベースのHTTPゲートウェイ。JSONパース、トークナイズ、バッチ分割などCPU処理を担い、gRPCプロトコルに変換してTulipへ転送する。本番トラフィックでペイロードのサイズがばらつく問題を、大バッチを複数チャンクに分割してレプリカ間でロードバランシングすることで対処する。
• Tulip：Rust + tokio + tonicで構築したgRPCサーバー。スケジューリングとバッチング後にエンジンへディスパッチする。
• ROSE（Runtime-Optimized Serving Engine）：主にPythonで実装したモデル推論レイヤー。カーネル、レイヤー定義、CUDAグラフ管理を担い、Tulipにstep()関数を公開する。ROSEはPerplexityが独自に開発したオープンソース推論エンジンであり、その詳細はPerplexityの紹介記事で解説されている。

重要な設計判断として、Perplexityはエンベディング専用エンジンを別途構築しなかった。エンベディングモデルは小型のTransformerであり、バッチエンベディング（ベクターDB構築時の高スループット処理）はLLMのprefillに、オンラインエンベディング（クエリ時の低レイテンシ処理）はdecodeにそれぞれ類似している。そのため、LLMスタックのprefill/decodeカーネルをそのまま流用している。

スケジューラーをあえてシンプルにした理由

Tulipのスケジューラーは先着順（FCFS）で、複雑な優先制御は持たない。これには測定に基づく根拠がある。Perplexityが扱うシーケンス長の小型エンベディングモデルでは、Attentionの二乗コストより全結合層の線形コストが支配的になる。つまりレイテンシはシーケンス数ではなくトークン数にほぼ比例する。サブ10億パラメータのモデルでは約512トークンでGPUが飽和し、それ以上バッチに詰め込んでも効率は上がらない。

CUDAグラフとLazyTensor：起動オーバーヘッドをどう削るか

小バッチではCPU側のカーネル起動コストがGPU実行コストを上回ることがある。Perplexityはすべてのエンベディングモデルに対してモデル全体のCUDAグラフを構築し、全カーネル起動を単一ドライバー呼び出しに畳み込んでいる。

ただし、FlashInferの一部のAttention実装がホスト側の動的入力に依存しており、全体グラフのキャプチャを阻害していた。この問題はPerplexityがFlashInferへのパッチをアップストリームすることで解決した。

グラフはトークン数の設定ごとにキャプチャが必要なため、トークン数を64または256の倍数にパディングしている。これにより数千のグラフが生成され、起動時に数分のキャプチャ時間が生じる。この問題をレイジーキャプチャで解消している。各設定で最初のリクエストはeagerに実行、2回目のヒット時にキャプチャとリプレイを起動する。起動直後のp99レイテンシは犠牲になるが、数分の初期化コストを数時間にわたって分散できる。

もう一つの工夫が**LazyTensor**だ。ページロック済みホストバッファとcudaMemcpyAsync、CUDAイベントを束ねた抽象で、step()が結果を同期待ちするのではなくLazyTensorを即座に返す。これによりRustの非同期タスクがバッチNの完了を待つ間、CPUがバッチN+1をエンキューできる。

カーネルバックエンドの選択

ROSEはragged入力（シーケンス長がバラバラなバッチ）に対して複数のAttentionバックエンドをサポートしている：FlashInfer 2、FlashInfer 3、FlashAttention 4。Perplexityによると一般的にはFlashAttention 4が高速だが、Qwen系モデルの非常に長いシーケンス長ではFlashInfer 3が上回るケースがあるため、バックエンド選択はモデルごとに判断している。なお、エンベディングモデル推論時はKVキャッシュをインスタンス化せず、パディングを避けるためragged attentionを使用する。

ベンチマーク比較

比較対象はvLLM v0.22.0で、BF16、実モデルウェイト、評価由来の入力を使用。コサイン類似度の乖離が0.1%以内であることをウォームアップで確認した上で、4種類のスイートを計測している。

• 低レイテンシエンベディング：バッチ1、128/512/4096トークン
• 低レイテンシスコアリング：バッチ5/25/50、512トークン
• 高スループットエンベディング：バッチ100、4並列プロセス
• 高並列エンベディング：並列1〜16、Ivyのトークナイズとネットワークオーバーヘッド込み

元記事によると、ROSEはvLLM v0.22.0に対してすべてのスイートで優位な結果を示しており、特に低レイテンシ・高並列条件での差が顕著だったとされている。具体的な数値やグラフは元記事のFast Embeddings on GPUsに掲載されているため、詳細はそちらを参照されたい。

なお、Ivy・Tulip・ROSEはPerplexity内部のコンポーネントであり、pplx-embedはPerplexityのEmbeddings API経由で外部からアクセスできる。

詳細はPerplexity Details Its GPU Embedding Stack: How Ivy, Tulip and ROSE Serve pplx-embedを参照していただきたい。