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 能力。