GPU 上的快速嵌入

快速且準確的搜尋對整個 Perplexity 至關重要，從搜尋和 Computer 到我們的 API 平台。在幕後，繁重的工作是由嵌入和排序模型完成的，這些模型可協助我們的系統針對給定的查詢識別出最相關的結果。透過訓練和服務我們自己的模型（例如 pplx-embed），我們實現了最先進的品質和延遲。

本文深入探討 Perplexity 針對此特殊類別模型之服務基礎架構的內部運作。我們討論了有效滿足 AI 原生搜尋推論需求的技術，實現模型的快速原型製作與評估，同時驅動我們的 EB 級搜尋索引。這些技術共同擴展了搜尋品質與效率的 Pareto 邊界，讓我們能以最低的成本和延遲為代理程式與使用者提供最佳的結果。

在典型的搜尋設定中，已建立索引的文件會使用嵌入模型對應至高維度向量空間，並儲存在向量資料庫中。透過使用相同的模型來嵌入查詢，可以藉由尋找最接近查詢向量的向量來定位相似的文件。這會為推論引擎衍生出兩種不同的流量模式：

批次嵌入：在建構、擴充或重新索引資料庫時，必須將大量文件嵌入到向量空間中，藉此將產出量最大化以將成本降至最低。

在進行向量搜尋之後，必須對大量文件進行評分，在產出量與延遲之間取得平衡。

線上嵌入：查詢資料庫時，必須內嵌一個簡短的查詢以便進行查找，從而將延遲降至最低。

我們建構了推論基礎架構，以便在各個使用情境中盡可能利用共用元件。由於我們通常使用小型 Transformer 模型來產生嵌入，因此我們與 LLM 推論程式碼共用了大部分的實作：批次嵌入類似於受限於計算能力的預填充（prefill），而通常在少數 token 上執行的線上嵌入則在計算上類似於受限於記憶體的解碼（decode）。因此，我們重複使用最佳化的預填充與解碼核心來提供嵌入模型服務。因此，我們能夠以極少的額外工程工作實現龐大的批次推論產出量，同時為線上嵌入工作負載保持低延遲。

我們透過標準化 API 公開推論，包括內部使用以及透過我們的 API 平台對外公開。在幕後，處理嵌入請求會涉及多個服務：

Ivy 是 Perplexity 服務會呼叫的 Rust HTTP 閘道。

它處理請求的 CPU 端工作，例如 JSON 解析、分詞、輸入範本化和批次分割，並將請求轉譯為自訂的 gRPC 通訊協定供下游伺服器使用。這種分離讓我們可以設定有關分詞和輸入格式的特定參數，而不必觸碰較重的推論實例。

Tulip 是推論伺服器介面。

這是一個使用 Rust、tokio 和 `tonic` 實作的 gRPC 伺服器。Tulip 接收 gRPC 推論請求，並處理排程與批次處理。接著，它會將批次傳送至 ROSE 引擎，並將完成的迴應傳回給客戶端。

**ROSE** (Runtime-Optimized Serving Engine) 實現了模型推論。

它主要以 Python 定義，為各種模型提供核心（kernels）、層與定義。ROSE 實作了模型的向前傳遞（forward passes），同時提供專門針對嵌入的 CUDA 圖形管理。它透過一個 step() 函式與 Tulip 橋接，該函式接收一個批次並傳回它在加速器上執行的運算參考。

基於 Transformer 的模型與底層的 Hopper/Blackwell 架構皆已是成熟技術，因此 GPU 端的嵌入推論在各個推論引擎中已收斂至大致最佳化的實作。即便如此，我們仍在執行階段與框架中發現了進一步改善的機會，能將模型端到端地暴露給客戶端。具體而言，我們發現透過謹慎管理 CUDA 圖形（CUDA graphs）並建構 LazyTensor 抽象層以非同步追蹤原生 Rust 引擎中的 GPU 端結果，可以改善延遲。我們在 Tulip 中實作了這些功能，使其能與 ROSE 的模型實作有效介接。

我們將 Tulip 設計成模型服務上盡可能輕量化的介面。它在 Tokio 非同步任務中處理傳入的請求，維護一個它所追蹤的請求集區，並從中排程批次以分派至加速器。Tulip 中的排程機制非常簡單：當 Tulip 正在分派工作或等待結果時，請求會不斷累積。從累積的請求中，系統會以先到先服務（first-come, first-served）的原則挑選序列來透過模型執行。

這項簡單的排程機制是由對模型效能的觀察所驅動。對於小型嵌入模型，在我們提供的序列長度下，我們注意到密集層的線性成本高於注意力的二次成本。因此，延遲主要與 token 數量成正比，而不是與序列數量成正比。因此，一旦批次大到足以使 GPU 飽和（在小於十億參數的模型上約為 512 個 token），在其中塞入更多序列並不會提高效率。

為了與模型進行有效介接，Tulip 依賴 CUDA 圖形與延遲結果追蹤來重疊 GPU 與 CPU 的工作，並充分利用可用資源。

執行模型的向前傳遞涉及 CPU 端與 GPU 端的工作。CPU 負責排程批次並以適當的參數啟動核心，而 GPU 則執行相關的矩陣乘法、注意力、正規化或啟用核心。對於訓練和重新索引等高產出量工作負載，由於批次大小和 GPU 端延遲都很大，因此 CPU 端的負荷微不足道。然而，在較小的批次大小下，CPU 端的工作可能會超過 GPU 端的工作。

為了減輕負荷，可以建構 CUDA 圖形來擷取透過單一呼叫 CUDA 驅動程式即可啟動向前傳遞之所有核心所需的中繼資料，而不是啟動獨立的核心。這消除了針對可以擷取 CUDA 圖形的組態重新執行昂貴的 Python 和 PyTorch 程式碼的需求。

在每個模型中，我們都會追蹤一個轉折點，以決定 GPU 執行成本高於 CPU 端核心啟動成本的最小 token 數量。由於嵌入模型很小，我們觀察到這個轉折點出現在數千個 token 和數十個序列的批次中。某些注意力機制實作依賴動態主機端輸入來設定核心啟動，因而無法使用完整模型的預填充/密集 CUDA 圖形。我們將相關核心的變更回饋至上游（upstreamed），以便在我們的推論引擎中啟用它們。

為了處理負荷，我們為所有嵌入模型建構了完整模型的 CUDA 圖形，並將 CPU 工作與 GPU 工作重疊。由於 CUDA 圖形將 CPU 端的負荷降到最低，因此一旦圖形啟動後，我們就有空閒時間在下一個批次可用時立即啟動並將其執行排入佇列。待處理批次的結果會透過 LazyTensor 進行追蹤，這允許 Rust 中的非同步任務進行封鎖，直到前一個批次執行完畢為止。CUDA 圖形有助於低延遲服務，確保我們不會受到核心啟動成本的拖累，並透過釋放 CPU 提早處理下一個批次的工作，促進高產出量情境下的更好排程。

必須針對每個不同的組態擷取 CUDA 圖形，對嵌入而言，這意味著每個序列計數與 token 計數組合都需要一個圖形。由於這個網格非常龐大，我們會將 token 計數填充至 64 或 256 倍數的區塊（buckets）。這仍然會產生數千個圖形，對於典型模型而言，可能需要花費數分鐘才能擷取完成。擷取成本來自兩個來源：必須執行以編譯核心並為需要它們的各種核心設定緩衝區的 eager 向前傳遞，接著是重新執行 Python 程式碼的擷取執行。

我們透過在引擎提供服務時以延遲方式擷取 CUDA 圖形，來減輕啟動成本。我們會追蹤每個組態，並確保它在觸發圖形擷取以及在第二次命中時重新播放之前，先經過 eager 暖機執行。隨後對相同圖形組態的所有執行都會透過 CUDA 圖形重新播放進行。延遲圖形擷取會對啟動期間的 p99 延遲造成影響；然而，它在將數分鐘的 eager 工作分散到多個小時方面非常有用。更快的啟動時間讓我們能更好地擴展和管理嵌入部署。

透過 CUDA，GPU 工作是非同步的。由於以非同步方式啟動核心會將其排入串流佇列中，因此主機程式碼必須明確進行同步才能讀出產生的向量。為了提升並行度，並能夠在等待前一個裝置端工作完成的同時啟動未來的批次，我們依賴 LazyTensor 抽象層來追蹤數值。

LazyTensor 透過事件追蹤頁面鎖定記憶體（page-locked memory）中的主機緩衝區以及複製裝置資料的 cudaMemcpyAsync 作業。它在同一個串流上啟動向前傳遞後觸發。由於複製作業必須等待串流上的所有先前核心執行完畢，因此相關事件會同時追蹤向前傳遞的完成情況以及 CPU 上結果的可用性。

我們在 ROSE 編碼器引擎中利用 LazyTensor 來重疊 GPU 與 CPU 的工作。每一個 step() 呼叫不再執行 CUDA 圖形並等待其完成，而是由 step() 傳回一個 LazyTensor 來非同步追蹤其結果。結合 CUDA 圖形，這有助於我們實現低延遲與更好的產出量。

我們調整了原本為 LLM 服務所建構的 ROSE 引擎，使其也能處理嵌入模型的執行。為了將支援嵌入模型所需的力氣降到最低，ROSE 在 LLM 與嵌入之間積極重複使用程式碼。例如，pplx-embed 服務與 Qwen3.5 LLM 解碼全部透過相同的核心進行。這種共用讓我們能輕鬆提供原本從 LLM 微調而來的嵌入模型，用於原型製作、評估與生產環境推論。

對於密集層（dense layers），由於 token 向量是獨立處理的，因此嵌入與 LLM 推論是相同的。在注意力層中，差異的處理方式是加入對不規則輸入（ragged inputs）的支援，以及 LLM 所需的分頁預填充與解碼設定。在提供嵌入模型服務時，我們不會實例化 KV 快取，而是分派至支援不規則格式的各種注意力核心，以避免填充。支援的轉換和校準常式也與 LLM 共用。

我們的推論 HTTP 代理層 Ivy 在效能上也扮演著重要角色。由於生產環境中的請求負載各不相同，將個別請求路由至個別複本可能會導致負載不平衡。Ivy 會將大批次請求分割成多個區塊，並在複本之間進行負載平衡，藉此提高利用率並平滑延遲。我們最近在 Ivy 中全面推出的自家 unigram 分詞（unigram tokenization）研究，大幅改善了現成分詞器的延遲。

ROSE 支援多種注意力後端。不同的核心可能適用於特定的問題大小。隨著時間推移，我們整合了 FlashInfer 2、FlashInfer 3 與 FlashAttention 4 核心來實作不規則注意力（ragged attention）。

總體而言，我們觀察到 FlashAttention 4 的速度更快。然而，在序列長度非常長的情況下，FlashInfer 3 在基於 Qwen 的模型上表現更佳。由於效能和調整可能會隨注意力機制的標頭數量與維度而異，因此我們在提供服務時會支援多種組態，並按情況逐一做出決定。

我們針對 vLLM v0.22.0 進行基準測試，在實際模型權重與衍生自評估資料集的輸入上以 BF16 精度執行推論。所有計時執行前均進行了暖機執行，以驗證餘弦相似度的差異在 0.1% 以內。

我們回報預先分詞的請求批次大小為 1、完全循序請求、序列長度分別為 128、512 和 4096 個 token 的執行時間。

預先分詞的請求批次大小為 5、25 和 50，序列長度為 512 個 token。

請求批次大小為 100，四個並行程序提交請求，序列長度分別為 512、1024 和 4096 個 token。

序列長度為 512，批次大小為 1，但我們發送 1、2、4、8 和 16 個並行請求。此基準測試也包含了透過 Ivy 進行的分詞成本，以及 Ivy 與 Tulip 之間的網路負荷。

由 Ivy、Tulip 和 ROSE 組成的服務基礎架構讓我們能以更低的延遲和更好的產出量為 Perplexity 提供嵌入服務，相較於現成解決方案，能以更低的成本實現更準確的搜尋。

透過專注於特定模型並掌控整個堆疊，我們獲得了在效能與靈活性之間取得有效平衡所需的自由，將高度可重複使用且高效能的 Rust 原語與更通用的 Python 建模程式碼混合使用。許多開源推論引擎（例如 vLLM、SGLang 與 TokenSpeed) 正在將 Rust 和 C++ 等語言整合到其堆疊中。在過去兩年中，我們投資於 Rust，並在效能和可維護性方面獲得了巨大回報。透過與我們的 LLM 服務堆疊共用大部分的嵌入實作，我們也提升了產出量，同時不需要花費大量的工程心力來維護嵌入模型。

隨著模型的演進，我們將持續改善堆疊的每一個層級，以減少受限於 CPU 與 GPU 的延遲。我們在 Ivy 與 Tulip 內部的自訂 gRPC 通訊協定讓我們能夠微調通訊以減少網路延遲，而 ROSE 則為提升計算產出量奠定了基礎。此外，隨著生態系統對無鎖執行緒 Python（free-threaded Python）支援的成長，我們將能進一步改善 Python-Rust 的互通性，以減少負荷。

搜尋嵌入

Tulips、Roses 和一些 Ivy

關注核心以外的事物

Tulip

CUDA 圖形管理

延遲張量

ROSE

Ivy

...但核心依然至關重要

基準測試

低延遲嵌入 (p50 / p90 / p99 / 最大值 ms)

低延遲評分 (p50 / p90 / p99 / 最大值 ms)

高產出量嵌入 (emb/s)

高並行嵌入 (p50 / p90 / p99 / 最大值 ms)

結論與未來工作

參考資料