GPU上の高速埋め込み
検索やコンピュータからAPIプラットフォームに至るまで、Perplexityのすべてにおいて、高速で正確な検索は不可欠です。裏では、埋め込みおよびランキングモデルが重労働を担っており、システムが特定のクエリに対して最も関連性の高い結果を特定するのに役立ちます。最先端の
検索やコンピュータからAPIプラットフォームに至るまで、Perplexityのすべてにおいて、高速で正確な検索は不可欠です。裏では、埋め込みおよびランキングモデルが重労働を担っており、システムが特定のクエリに対して最も関連性の高い結果を特定するのに役立ちます。pplx-embedなどの独自のモデルをトレーニングしてサービングすることで、最先端の品質とレイテンシを実現しています。
この記事では、この特殊なクラスのモデルに対するPerplexityのサービングインフラストラクチャの内部構造を紹介します。AIネイティブ検索の推論ニーズに効率的に対応し、モデルの迅速なプロトタイピングと評価を可能にしながら、exabyte規模の検索インデックスを動かすための手法について議論します。これらの手法は総合的に、検索の品質と効率のパレートフロンティアを拡大し、最小限のコストとレイテンシでエージェントとユーザーに最高の成果を提供することを可能にします。
検索のための埋め込み
通常の検索セットアップでは、インデックス化されたドキュメントは埋め込みモデルを使用して高次元ベクトルスペースにマッピングされ、ベクトルデータベースに格納されます。同じモデルを使用してクエリを埋め込むことで、クエリのベクトルに最も近いベクトルを見つけることによって、類似したドキュメントを特定できます。これにより、推論エンジンが処理すべき2つの異なるトラフィックパターンが生じます。
- バッチ埋め込み：データベースの構築、拡張、または再インデックス作成時に、一括ドキュメントをベクトルスペースに埋め込み、スループットを最大化してコストを最小限に抑えます。
ベクトル検索に続いて、大量のドキュメントのスコアリングを行う必要があり、スループットとレイテンシのバランスを取ります。
- オンライン埋め込み：データベースをクエリする際、ルックアップのために短いクエリを埋め込み、レイテンシを最小限に抑えます。
ユースケース間で可能な限り共通のコンポーネントを活用できるように推論インフラストラクチャを構築しました。埋め込みの生成には通常小さなTransformerモデルを使用するため、実装の大部分をLLM推論コードと共有しています。バッチ埋め込みは計算バウンドなプリフィルに似ており、一方、少数のトークンで実行されることが多いオンライン埋め込みは、メモリバウンドなデコードに計算上似ています。したがって、最適化されたプリフィルおよびデコードカーネルを再利用して埋め込みモデルを提供しています。その結果、オンライン埋め込みワークロードの低レイテンシを維持しながら、追加のエンジニアリング作業を最小限に抑えて大規模なバッチ推論スループットを実現できます。
Tulips、Roses、そしていくつかのIvy
内部およびAPIプラットフォームを通した外部の両方で、標準化されたAPIを通じて推論を公開します。内部では、埋め込みリクエストの処理に複数のサービスが関与しています。
- Ivyは、Perplexityのサービスが呼び出すRustのHTTPゲートウェイです。
JSON解析、トークン化、入力のテンプレート化、バッチ分割などのリクエストに対するCPU側の作業を処理し、リクエストを下流サーバー用のカスタムgRPCプロトコルに変換します。この分離により、より重い推論インスタンスに触れることなく、トークン化や入力フォーマットに関する特定のパラメータを構成できます。
- Tulipは推論サーバーインターフェイスです。
Rust、tokio、`tonic`で実装されたgRPCサーバーです。TulipはgRPC推論リクエストを受信し、スケジューリングとバッチ処理を処理します。その後、バッチをROSEエンジンに送信し、完了したレスポンスをクライアントに返します。
- **ROSE**（Runtime-Optimized Serving Engine）はモデル推論を実装します。
主にPythonで定義されており、さまざまなモデル向けにカーネル、レイヤー、定義を提供します。ROSEはモデルを通したフォワードパスを実装し、埋め込みに特化したCUDAグラフ管理も提供します。バッチを受け取り、アクセラレータ上で実行される計算への参照を返すstep()関数を介してTulipにブリッジされます。
カーネルを超えて注意を払う
Transformerベースのモデルと、その下層にあるHopper/Blackwellアーキテクチャはいずれも成熟したテクノロジーであるため、GPU側での埋め込み推論は、さまざまな推論エンジンの間で概ね最適な実装へと収束しています。それでも、モデルをクライアントにエンドツーエンドで公開するランタイムやハーネスには、さらなる改善の余地があることがわかりました。特に、CUDAグラフを慎重に管理し、ネイティブのRustエンジン内でGPU側の結果を非同期に追跡するLazyTensor抽象化を構築することで、レイテンシを改善できることが判明しました。これらの機能をTulipに実装することで、ROSEのモデル実装と効果的に連携できるようになりました。
Tulip
Tulipは、モデルサービングに対するインターフェイスとして可能な限り軽量になるように設計されています。Tokioの非同期タスクで受信リクエストを処理し、追跡するリクエストのプールを維持し、アクセラレータにディスパッチするバッチをスケジュールします。Tulipのスケジューリングメカニズムは非常にシンプルです。Tulipが作業をディスパッチしているか結果を待っている間に、リクエストが蓄積されます。蓄積されたリクエストから、先着順でシーケンスが選ばれ、モデルを通過します。
シンプルなスケジューリングメカニズムは、モデルパフォーマンスに関する観察に基づいています。小さな埋め込みモデルの場合、サービングするシーケンス長において、高密度レイヤーの線形コストがアテンションの二次コストよりも支配的であることがわかりました。したがって、レイテンシはシーケンス数ではなく、主にトークン数に比例します。その結果、バッチがGPUを飽和させるのに十分な大きさ（10億パラメータ未満のモデルで約512トークン）になると、さらに多くのシーケンスを詰め込んでも効率は向上しません。
モデルと効果的に連携するために、TulipはCUDAグラフと遅延結果追跡に依存して、GPUとCPUの作業をオーバーラップさせ、利用可能なリソースを完全に活用します。
CUDAグラフ管理
モデルのフォワードパスの実行には、CPU側とGPU側の両方の作業が含まれます。CPUはバッチのスケジュールと適切なパラメータでのカーネルの起動を担当し、GPUは関連する行列乗算、アテンション、正規化、または活性化カーネルを実行します。トレーニングや再インデックス作成などの高スループットのワークロードでは、バッチサイズとGPU側のレイテンシがいずれも大きいため、CPU側のオーバーヘッドは無視できるほど小さくなります。しかし、より小さなバッチサイズでは、CPU側の作業がGPU側の作業を上回る可能性があります。
オーバーヘッドを軽減するため、個別のカーネルを起動する代わりに、CUDAドライバへの単一の呼び出しでフォワードパスのすべてのカーネルを起動するために必要なメタデータをキャプチャするCUDAグラフを構築できます。これにより、CUDAグラフをキャプチャできる構成のために、コストのかかるPythonおよびPyTorchコードを再実行する必要がなくなります。
各モデルにおいて、GPU実行がCPU側のカーネル起動よりもコストが高くなる最小限のトークン数を特定し、変曲点を追跡しています。埋め込みモデルは小さいため、この変曲点は数千トークンおよび数十シーケンスのバッチで現れることが観察されます。一部のアテンション実装は、カーネルの起動を構成するために動的なホスト側の入力に依存しており、フルモデルのプリフィル/高密度CUDAグラフを防いでいます。推論エンジンでこれらを有効にするため、関連するカーネルへの変更をアップストリームに提供しました。
オーバーヘッドに対処するため、すべての埋め込みモデルに対してモデル全体のCUDAグラフを構築し、CPUの作業とGPUの作業をオーバーラップさせます。CUDAグラフはCPU側のオーバーヘッドを最小限に抑えるため、グラフが起動されると、次のバッチが利用可能になったときにいつでもその実行を開始してキューに入れるための空き時間が生まれます。保留中のバッチの結果はLazyTensorで追跡され、これによりRustの非同期タスクは前のバッチの実行が完了するまでブロックできるようになります。CUDAグラフは、カーネル起動のコストに妨げられないようにすることで低レイテンシサービングを支援し、CPUを解放して次のバッチの作業をより早く行えるようにすることで、高スループットケースでのスケジューリングの改善を促進します。
CUDAグラフは個別の構成ごとにキャプチャする必要があります。埋め込みの場合、これはシーケンス数とトークン数の組み合わせごとにグラフが存在することを意味します。このグリッドは広範囲にわたるため、トークン数を64または256の倍数のバケットにパディングします。これにより、一般的なモデルではキャプチャに数分かかる可能性のある数千のグラフが依然として生成されます。キャプチャのコストは2つのソースから発生します。カーネルをコンパイルし、それを必要とするさまざまなカーネル用のバッファを設定するために実行する必要がある急欲なフォワードパスと、それに続くPythonコードを再実行するキャプチャ実行です。
エンジンがサービングするにつれてCUDAグラフを遅延キャプチャすることで、起動コストを軽減します。各構成を追跡し、2回目のヒットでグラフのキャプチャと再生をトリガーする前に、急欲なウォームアップ実行を確実に通過させます。同じグラフ構成の以降のすべての実行は、CUDAグラフの再生を経ます。遅延グラフキャプチャは、起動時のp99レイテンシに影響を与えますが、数分の急欲な作業を複数時間にわたって分散させる上で価値があります。起動時間が短縮されることで、埋め込みデプロイメントのスケールと管理が容易になります。
遅延テンソル
CUDAを通じて、GPUの作業は非同期です。カーネルを非同期で起動するとストリームにキューイングされるため、ホストコードは結果のベクトルを読み出すために明示的に同期する必要があります。より高いレベルの並列処理を容易にし、デバイス上での前のバッチの完了を待ちながら将来のバッチを開始できるようにするために、値を追跡するためのLazyTensor抽象化に依存しています。
LazyTensorは、ページロックメモリ内のホストバッファと、デバイスからデータをコピーするイベントを介したcudaMemcpyAsync操作を追跡します。これは、同じストリームでのフォワードパスの起動後に開始されます。コピー操作はストリーム上の以前のすべてのカーネルの実行を待機する必要があるため、関連するイベントはフォワードパスの完了とCPU上の結果の利用可能性の両方を追跡します。
GPUとCPUの作業をオーバーラップさせるため、ROSEエンコーダーエンジンでLazyTensorを活用しています。各step()呼び出しがCUDAグラフを実行してその完了を待つ代わりに、step()は結果を非同期に追跡するLazyTensorを返します。CUDAグラフと組み合わせることで、低レイテンシと優れたスループットの達成に貢献します。
ROSE
元々はLLMサービング用に構築したROSEエンジンを適応させ、埋め込みモデルの実行も処理できるようにしました。埋め込みモデルをサポートするために必要な労力を最小限に抑えるため、ROSEはLLMと埋め込みの間でコードを積極的に再利用します。たとえば、pplx-embedサービングとQwen3.5のLLMデコードはすべて同じカーネルを通過します。この共有により、プロトタイピング、評価、本番推論のために、元々LLMからファインチューニングされた埋め込みモデルを簡単に提供できるようになります。
高密度レイヤーの場合、トークンベクトルが独立して処理されるため、埋め込みとLLM推論は同一です。アテンションレイヤーでは、LLMに必要なページ化されたプリフィルおよびデコード設定に加えて、ラギッド入力のサポートを追加することで違いが処理されます。埋め込みモデルを提供する場合、KVキャッシュをインスタンス化せず、パディングを避けるためにラギッド形式をサポートするアテンションのバリエーションにディスパッチします。サポートする変換および調整ルーチンもLLMと共有されます。
Ivy
推論HTTPプロキシレイヤーであるIvyも、パフォーマンスにおいて重要な役割を果たしています。リクエストのペイロードは本番環境によって異なるため、個別のリクエストを個別のレプリカにルーティングすると、負荷の不均衡を引き起こす可能性があります。Ivyは大規模バッチリクエストをチャンクに分割し、レプリカ間で負荷分散を行うことで、利用率を向上させ、レイテンシを平滑化します。Ivyで完全にロールアウトされた自社製のユニグラムトークン化に関する最近の作業により、市販のトークナイザーと比較してレイテンシが大幅に改善されました。
...しかし、それでもカーネルは重要である
ROSEはさまざまなアテンションバックエンドをサポートしています。異なるカーネルが特定の問題サイズに適している場合があります。時間の経過とともに、FlashInfer 2、FlashInfer 3、FlashAttention 4の各カーネルを統合して、ラギッドアテンションを実装しました。
一般に、FlashAttention 4の方が高速であることが観測されています。ただし、非常に長いシーケンス長を持つQwenベースのモデルでは、FlashInfer 3の方が優れたパフォーマンスを発揮します。パフォーマンスとチューニングはアテンションヘッドの数や次元によって異なるため、複数の構成に対するサポートを維持し、サービング時にケースバイケースで判断を下しています。
ベンチマーク
評価データセットから得られた実際のモデルウェイトと入力に対し、BF16精度で推論を実行するvLLM v0.22.0と比較してベンチマークを実施しています。すべてのタイミング実行の前にウォームアップ実行が行われ、コサイン類似度の乖離が0.1%以内であることが検証されています。
低レイテンシ埋め込み（p50 / p90 / p99 / 最大 ms）
事前トークン化されたリクエストバッチサイズ1、完全にシーケンシャルなリクエスト、シーケンス長128、512、4096トークンのランタイムを報告します。
低レイテンシスコアリング（p50 / p90 / p99 / 最大 ms）
事前トークン化されたリクエストのバッチサイズは5、25、50、シーケンス長は512トークン。
高スループット埋め込み（emb/s）
リクエストバッチサイズ100、リクエストを送信する4つの並行プロセス、シーケンス長は512、1024、4096トークン。
高並行性埋め込み（p50 / p90 / p99 / 最大 ms）
シーケンス長512、バッチサイズ1ですが、1、2、4、8、16の並行リクエストを送信します。このベンチマークには、Ivyを介したトークン化コストと、IvyとTulip間のネットワークオーバーヘッドも含まれます。
結論と今後の展望
Ivy、Tulip、ROSEで構成されるサービングインフラストラクチャにより、Perplexityの埋め込みをより低いレイテンシと優れたスループットで提供できるようになり、市販のソリューションと比較して、削減されたコストでより正確な検索結果をもたらします。
特定のモデルに焦点を当ててスタック全体を所有することにより、パフォーマンスと柔軟性の効果的なバランスを取るために必要な自由度が得られ、高度に再利用可能でパフォーマンスの高いRustプリミティブと、より汎用的なPythonモデリングコードを混在させています。vLLM、SGLang、TokenSpeedなどの多くのオープンソース推論エンジンは、RustやC++などの言語をスタックに統合しています。当社は過去2年間にわたりRustに投資し、パフォーマンスと保守性の両方で大きな成果を上げてきました。埋め込み実装の大部分をLLMサービングスタックと共有することで、埋め込みモデルのメンテナンスに多大なエンジニアリング労力を費やすことなく、スループットの向上も実現しています。
モデルの進化に伴い、CPUバウンドとGPUバウンドの両方のレイテンシを削減するために、スタックの各層の改善を続けます。IvyとTulip内のカスタムgRPCベースのプロトコルにより、通信を微調整してネットワークレイテンシを削減できる一方、ROSEは計算スループットを向上させるための基盤を提供します。さらに、エコシステム全体でフリーレッドPythonのサポートが成長するにつれて、PythonとRustの相互運用性をさらに改善してオーバーヘッドを削減できるようになります。