GPU 上的快速嵌入

快速且準確的搜尋對整個 Perplexity 至關重要,從「搜尋和電腦」(Search and Computer)到我們的 API 平台莫不如此。幕後的核心工作是由嵌入和排序模型完成的,它們幫助我們的系統為給定的查詢識別出最相關的結果。我們實現了最先進的

作者Perplexity Engineering

快速且準確的搜尋對整個 Perplexity 至關重要,從「搜尋和電腦」(Search and Computer)到我們的 API 平台莫不如此。幕後的核心工作是由嵌入和排序模型完成的,它們幫助我們的系統為給定的查詢識別出最相關的結果。透過訓練和提供我們自己的模型(例如 pplx-embed),我們實現了最先進的品質和延遲。

本文深入探討了 Perplexity 針對此類特殊模型的服務基礎設施幕後運作。我們討論了有效滿足 AI 原生搜尋推論需求的技術,能夠快速進行模型原型設計與評估,同時驅動我們的 艾位元組規模搜尋索引(exabyte-scale search index)。這些技術共同擴展了搜尋品質與效率的帕累托前沿(Pareto frontier),使我們能夠以最低的成本和延遲,為代理和使用者提供最佳的結果。

用於搜尋的嵌入

在典型的搜尋設定中,已索引的文件會使用嵌入模型對應到高維度向量空間,並儲存在向量資料庫中。透過使用相同的模型來嵌入查詢,可以透過尋找最接近該查詢向量的向量來定位相似的文件。這為推論引擎帶來了兩種不同的流量模式來提供服務:

  • 批次嵌入(Batch Embedding):在建構、擴充或重新索引資料庫時,必須將大量文件嵌入到向量空間中,以將吞吐量最大化並將成本降到最低。

在向量搜尋之後,必須對大量文件進行評分,在吞吐量和延遲之間取得平衡。

  • 線上嵌入(Online Embedding):在查詢資料庫時,必須為查找嵌入一個短查詢,以將延遲降至最低。

我們建構了推論基礎設施,以便在各種使用情境中盡可能利用共用元件。由於我們通常使用小型 Transformer 模型來產生嵌入,因此我們與 LLM 推論程式碼共用了大部分的實現:批次嵌入類似於計算受限的預填充(prefill),而通常對少數 Token 執行的線上嵌入則在計算上類似於記憶體受限的解碼(decode)。因此,我們重複使用經過 最佳化的預填充與解碼核心 來提供嵌入模型服務。因此,我們可以在進行極少額外工程工作的情況下,實現巨大的批次推論吞吐量,同時為線上嵌入工作負載保持低延遲。

Tulips、Roses 與一些 Ivy

我們透過標準化 API 公開推論,包括內部使用以及透過我們的 API 平台對外提供。在幕後,有多個服務參與了嵌入請求的處理:

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

它處理請求的 CPU 端工作,例如 JSON 解析、分詞、輸入模版化(input templating)和批次分割,並將請求轉譯為用於下游伺服器的自訂 gRPC 協定。這種分離使我們能夠配置關於分詞和輸入格式化的特定參數,而不必觸及較重的推論實例。

  • Tulip 是推論伺服器介面。

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

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

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

從嵌入請求經過 Ivy 到複製的 Tulip 伺服器的服務架構

關注核心之外的層面

Transformer 架構模型和底層的 Hopper/Blackwell 架構都是成熟的技術,因此在 GPU 端進行嵌入推論時,各種推論引擎的實現方式已趨於大致最佳化。即便如此,我們仍在執行階段和封裝程式中發現了進一步改善的機會,能夠將模型端到端地暴露給客戶端。特別是,我們發現透過小心管理 CUDA 圖形(CUDA graphs),並建構 LazyTensor 抽象層在原生 Rust 引擎中非同步追蹤 GPU 端的結果,可以有效改善延遲。我們在 Tulip 中實現了這些功能,使其能與 ROSE 的模型實現進行有效介接。

Tulip

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

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

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

CUDA 圖形管理

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

急切向前傳遞(Eager forward pass):與裝置核心交錯的主機呼叫

為緩解負荷,建構 CUDA 圖形可以捕捉啟動向前傳遞所需的所有核心所需的後設資料(metadata),只需對 CUDA 驅動程式進行一次呼叫,而不必啟動獨立的核心。這消除了針對可以擷取 CUDA 圖形的配置重新執行昂貴的 Python 和 PyTorch 程式碼的需求。

針對每個模型,我們都會追蹤一個轉折點,以確定 GPU 執行成本高於 CPU 端核心啟動成本的最小 Token 數量。由於嵌入模型很小,我們發現這個轉折點出現在數千個 Token 和數十個序列的批次中。某些注意力機制實現依賴動態主機端輸入來配置核心啟動,因而妨礙了完整模型的預填充/密集 CUDA 圖形。我們將相關核心的變更上游化(upstreamed),以便在我們的推論引擎中啟用它們。

為了解決負荷問題,我們為所有嵌入模型建構了全模型 CUDA 圖形,並將 CPU 工作與 GPU 工作重疊。由於 CUDA 圖形將 CPU 端的負荷降到最低,一旦圖形啟動,我們就有空閒時間在下一個批次可用時立即啟動並將其執行排入佇列。待處理批次的結果會由 LazyTensor 進行追蹤,這允許 Rust 中的非同步任務進行阻塞,直到前一個批次執行完畢。CUDA 圖形有助於低延遲服務,因為它們確保我們不會被核心啟動的成本拖慢,並且透過釋放 CPU 更快地處理下一個批次的工作,促進了高吞吐量情況下的改進排程。

Cudagraph 向前傳遞

必須為每個不同的配置擷取 CUDA 圖形,對嵌入而言,這意味著每個序列計數和 Token 計數組合都需要一個圖形。由於這個網格非常龐大,我們會將 Token 計數填補(pad)到 64 或 256 的倍數的區塊(buckets)中。這仍然會產生數千個圖形,對於典型模型而言,可能需要數分鐘才能擷取完畢。擷取成本來自兩個來源:必須執行急切向前傳遞以編譯核心並為需要它們的各種核心設定緩衝區,接著是重新執行 Python 程式碼的擷取運行。

我們透過在引擎提供服務時以延遲方式擷取 CUDA 圖形來緩解啟動成本。我們會追蹤每個配置,並確保它在觸發圖形擷取以及在第二次命中時進行重播之前,先通過急切預熱執行(eager warmup run)。隨後對同一圖形配置的所有執行都將透過 CUDA 圖形重播進行。延遲圖形擷取會對啟動時的 p99 延遲產生影響;然而,它在將數分鐘的急切工作分散到多個小時方面非常有價值。更快的啟動時間讓我們能夠更好地擴展和管理嵌入部署。

延遲張量(Lazy Tensors)

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

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

LazyTensor 執行

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

顯示 CPU 準備與同步重疊連續 GPU 批次的時間軸

ROSE

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

對於密集層(dense layers),由於 Token 向量是獨立處理的,因此嵌入與 LLM 推論是相同的。在注意力層中,差異透過新增對不規則(ragged)輸入的支援來處理,同時輔以 LLM 所需的分頁預填充與解碼設定。當提供嵌入模型服務時,我們不會實例化 KV 快取(KV cache),而是分派至支援不規則格式的各種注意力機制核心,以避免填充(padding)。支援的轉換和校準常式也與 LLM 共用。

Ivy

Ivy 作為我們的推論 HTTP 代理層,在效能上也扮演著重要角色。由於請求酬載(payloads)在生產環境中會有所不同,將個別請求路由至個別副本可能會導致負載不平衡。Ivy 會將大批次請求分割成多個區塊(chunks),並在副本之間進行負載平衡,從而提高利用率並平滑延遲。我們近期在 自研一元分詞(in-house unigram tokenization) 方面的工作已在 Ivy 中全面推出,大幅改善了現成分詞器的延遲。

……但核心依然至關重要

ROSE 支援各種注意力後端(backends)。不同的核心可能適合特定的問題大小。隨著時間推移,我們整合了 FlashInfer 2、FlashInfer 3 和 FlashAttention 4 核心來實現不規則注意力(ragged attention)。

依模型形狀與問題大小的注意力核心效能

總體而言,我們觀察到 FlashAttention 4 的速度更快。然而,在極長的序列長度下,FlashInfer 3 在基於 Qwen 的模型上表現更佳。由於效能和調整可能會隨注意力機制的頭數與維度而變化,因此我們在提供服務時會支援多種配置,並依實際情況做出決定。

效能評測

我們以 vLLM v0.22.0 為基準進行評測,在 BF16 精度下對衍生自評估資料集的實際模型權重和輸入執行推論。所有的計時執行都在預熱執行之後進行,預熱執行已驗證餘弦相似度的差異在 0.1% 以內。

低延遲嵌入(p50 / p90 / p99 / 最大毫秒數)

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

低延遲嵌入效能評測結果

低延遲評分(p50 / p90 / p99 / 最大毫秒數)

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

低延遲評分效能評測結果

高吞吐量嵌入(emb/s)

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

高吞吐量嵌入效能評測結果

高併發嵌入(p50 / p90 / p99 / 最大毫秒數)

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

高併發嵌入效能評測結果

結論與未來工作

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

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

隨著模型的演進,我們將持續改善堆疊的每一層,以減少 CPU 受限和 GPU 受限的延遲。我們在 Ivy 和 Tulip 內部的自訂 gRPC 協定讓我們能夠微調通訊以減少網路延遲,而 ROSE 則為提高計算吞吐量提供了基礎。此外,隨著生態系統對無鎖執行緒 Python(free-threaded Python)的支援不斷增長,我們將能夠進一步改善 Python 與 Rust 的互通性,以減少負荷。

參考資料