Perplexityのエンジニアリングチームは「GPU上での高速埋め込み」と題する技術ブログを公開し、埋め込みモデルpplx-embedを支える推論基盤の詳細を明らかにした。AI検索製品の検索品質は、埋め込みモデルの性能と、それを大規模インデックスに対して低コストで運用できるかという2つの要素で決まる。今回の記事は後者、すなわちPerplexity Search、Computer、API Platformで使われるpplx-embedとランキングモデルのためのサービング基盤に焦点を当てている。
Perplexityは埋め込みサービスのトラフィックを2つの種類に整理している。バッチ埋め込みはベクトルデータベースの構築や再インデックス時に発生し、スループットがコストを左右する。オンライン埋め込みはクエリ時に発生し、短いクエリを高速にベクトル化する必要がある。スコアリングはその中間に位置し、ベクトル検索後に大量のドキュメントをランキングするため、両方の特性が求められる。
重要な設計判断として、Perplexityは専用の埋め込みエンジンを構築しなかった。埋め込みモデルは小型のTransformerであり、バッチ埋め込みは計算量が多いprefillに、数トークンのオンライン埋め込みはメモリ帯域に依存するdecodeに類似している。そのため、研究チームはLLMスタックのprefill用・decode用カーネルをそのまま再利用している。
リクエスト処理は3つのサービスで構成される。IvyはRust製のHTTPゲートウェイで、JSON解析、トークナイゼーション、入力テンプレート処理、バッチ分割などのCPU側処理を担当し、独自のgRPCプロトコルに変換する。また、大規模バッチを複数のチャンクに分割し、レプリカ間で負荷分散することで、本番環境で発生するペイロードサイズのばらつきによる不均衡を補正する。Tulipは推論サーバのインターフェースであり、Rust、tokio、tonicで構築されたgRPCサーバとして、エンジンへの投入前にスケジューリングとバッチ処理を行う。ROSEは推論エンジン本体で、主にPythonで実装され、カーネルやレイヤー、モデル定義を提供し、CUDA Graphを管理して、Tulipに対してstep()関数を公開する。
スケジューラは意図的にシンプルに保たれている。Tulipはリクエストを先着順で選択し、その間にバッチを蓄積する。この単純さは計測に基づいている。小型の埋め込みモデルでは、Perplexityが扱うシーケンス長の範囲で、密層の線形コストがアテンションの2次コストを上回る。したがって、レイテンシはシーケンス数ではなくトークン数にほぼ比例する。GPUを飽和させるバッチ（サブ10億パラメータモデルでは約512トークン）に達すると、追加のシーケンスを詰め込んでも効率は向上しない。
小バッチでは、CPU側のカーネル起動時間がGPU実行時間を上回ることがある。Perplexityは全埋め込みモデルに対してモデル全体のCUDA Graphを構築し、すべての起動を1回のドライバ呼び出しにまとめている。埋め込みモデルは小さいため、GPU作業が起動コストを上回る変曲点は、数千トークン・数十シーケンスのバッチサイズで現れる。一部のアテンション実装は動的なホスト側入力に依存しており、モデル全体のグラフ構築を妨げる。そのため、PerplexityはFlashInferに対して変更をアップストリームし、キャプチャを可能にした。
CUDA Graphは構成ごとにキャプチャする必要があり、トークン数は64または256の倍数にパディングされる。これでもモデルごとに数千のグラフが生成され、キャプチャに数分かかることがある。Perplexityは遅延キャプチャ方式を採用し、各構成で最初にeagerウォームアップを行い、2回目の出現時にキャプチャとリプレイをトリガーする。これにより起動時のp99レイテンシにコストがかかるが、数分間のeager作業を数時間に分散できる。
もう1つの要素はLazyTensorである。ページロックされたホストバッファ、cudaMemcpyAsync、CUDAイベントを追跡し、step()がデバイス上でブロックする代わりにLazyTensorを返すことで、Rustの非同期タスクがバッチNの完了を待っている間に、CPUはバッチN+1をエンキューできる。
カーネルレベルでも最適化が行われる。ROSEは非均一入力のための複数のアテンションバックエンド（FlashInfer 2、FlashInfer 3、FlashAttention 4）をサポートする。報告によると、FlashAttention 4は一般的に高速だが、Qwenベースのモデルで非常に長いシーケンス長の場合、FlashInfer 3が優れており、バックエンドはケースバイケースで選択される。埋め込みモデルを提供する際にはKVキャッシュを生成せず、パディングを避けるために非均一アテンションのバリアントを利用する点も特徴的である。
ベンチマークでは、vLLM v0.22.0とBF16精度で比較し、実重みと評価由来の入力を使用して、ウォームアップ実行でコサイン類似度の乖離が0.1%以内であることを確認している。低レイテンシ埋め込み（バッチ1、128/512/4096トークン）、低レイテンシスコアリング（バッチ5/25/50、512トークン）、高スループット埋め込み（バッチ100、4並行プロセス）、高並列埋め込み（1〜16並行リクエスト、Ivyのトークナイゼーションとネットワークオーバーヘッドを含む）の4つのスイートで評価された。
結論として、Perplexityの埋め込みスタックは、LLM用の推論カーネルを再利用し、Ivy、Tulip、ROSEという3つの内部サービスが連携することで、高い性能を実現している。CUDA Graphの遅延キャプチャやLazyTensorなどの工夫が、バッチサイズに応じたオーバーヘッドを効果的に削減している。これらの技術は内部的に利用されており、外部開発者はPerplexityのEmbeddings APIを通じてpplx-embedを利用できる。