Perplexity 工程團隊近日釋出了題為《GPU 上的快速嵌入》的技術部落格，首次詳細介紹了支撐其嵌入模型 pplx-embed 的推理基礎設施。在 AI 搜尋產品中，檢索質量主要受兩個因素制約：一是嵌入模型本身的質量，二是如何低成本地在海量索引上執行該模型。Perplexity 的這篇文章聚焦後者，說明了為 pplx-embed 以及搜尋、計算機和 API 平臺中的排序模型所構建的服務架構。
Perplexity 將嵌入服務劃分為兩種流量模式：批處理嵌入和線上嵌入。批處理嵌入發生在構建或重建向量資料庫時，此時吞吐量決定成本；線上嵌入發生在查詢時，短查詢需要極低的延遲。排序任務則介於兩者之間：向量檢索後需要為大批次文件打分，需要在吞吐和延遲之間取得平衡。
架構上的關鍵決策是：Perplexity 並沒有為嵌入服務單獨開發一個推理引擎。由於嵌入模型本質上是小型 Transformer，批處理嵌入類似於計算密集的 prefill 階段，而線上嵌入通常只包含幾個 token，類似於記憶體密集的 decode 階段。因此，開發團隊直接複用了其 LLM 推理棧中的 prefill 和 decode 核心。
整個請求鏈路由三個元件構成。Ivy 是一個 Rust 編寫的 HTTP 閘道器，負責 CPU 側的工作，包括 JSON 解析、分詞、輸入模板處理和批次切分，並將請求轉換為自定義的 gRPC 協議。它還能將大批次請求拆分為更小的塊，並在各個副本間進行負載均衡，從而解決生產環境中負載不均衡的問題。Tulip 是推理伺服器介面，使用 Rust、tokio 和 tonic 構建，負責任務排程和批次組裝，然後將其分發給引擎。ROSE 則是推理引擎，主要使用 Python，提供核心、層和模型定義，管理 CUDA 圖，並向 Tulip 暴露 step() 函式。
排程器的設計刻意保持簡單：Tulip 採用先到先服務的方式處理請求，同時積累批次。這種設計源於一個觀察：對於小尺寸嵌入模型，在 Perplexity 使用的序列長度範圍內，稠密層的線性計算成本往往高於注意力機制帶來的平方級成本。因此，延遲大致與 token 數成正比，而不是序列數。當批次規模足以讓 GPU 飽和（在亞十億引數模型上大約為 512 個 token）時，繼續增加序列數並不會帶來額外的效率收益。
為了減少小批次場景下的 CPU 端核心啟動開銷，Perplexity 為所有嵌入模型構建了整模型 CUDA 圖，將數十次核心啟動合併為一次驅動呼叫。由於嵌入模型較小，GPU 計算超過啟動開銷的臨界點通常出現在數千 token 和數十序列的批次規模。部分注意力實現依賴動態主機端輸入，會阻礙整圖捕獲；Perplexity 為此將改動上游提交給 FlashInfer，解決了該問題。CUDA 圖需要針對每種配置捕獲，因此 token 數會按 64 或 256 的倍數進行填充。這仍然可能產生數千張圖，每模型需要數分鐘捕獲時間。Perplexity 採用懶捕獲策略：每種配置先進行一次 eager 預熱，在第二次命中時才觸發捕獲和重放。這雖然會帶來啟動階段的 p99 延遲代價，但將數分鐘的 eager 工作分散到數小時內。
另一個關鍵元件是 LazyTensor。它跟蹤一個頁鎖定主機緩衝區，並關聯一次 cudaMemcpyAsync 操作和一個 CUDA 事件。step() 不會阻塞等待 GPU 完成，而是立即返回一個 LazyTensor，使得 Rust 非同步任務可以在等待批次 N 完成的同時，CPU 繼續入隊批次 N+1。這種設計實現了 CPU 端的批准備與 GPU 執行的重疊。
核心層面依然存在最佳化空間。ROSE 支援多種適用於非均勻輸入的注意力後端，包括 FlashInfer 2、FlashInfer 3 和 FlashAttention 4。Perplexity 的報告表明，FlashAttention 4 在大多數情況下更快，但在基於 Qwen 的超長序列模型中，FlashInfer 3 表現更好，因此後端選擇是逐個案例進行的。值得注意的是，當 ROSE 用於嵌入模型時，不會例項化 KV 快取，而是排程到非均勻注意力實現以避免填充。
在基準測試方面，Perplexity 與 vLLM v0.22.0 在 BF16 精度下進行了對比，使用真實權重和評估資料，並透過預熱執行驗證餘弦相似度差異不超過 0.1%。測試涵蓋四個場景：低延遲嵌入（batch 1，128/512/4096 token）、低延遲排序（batch 5/25/50，512 token）、高吞吐嵌入（batch 100，四個併發程序）以及高併發嵌入（1 到 16 個併發請求，包含 Ivy 的 tokenization 和網路開銷）。測試結果顯示，該架構在延遲與吞吐上均具有競爭力。
總體而言，Perplexity 的嵌入服務棧以複用 LLM 核心為核心，用三個內部服務協同完成請求處理，並透過 CUDA Graph、懶捕獲和 LazyTensor 實現了顯著的效能提升。這些元件目前屬於內部基礎設施，外部使用者可透過 Perplexity 的 Embeddings API 使用 pplx-embed 能力。