Perplexity AI 開源Unigram分詞器，p50延遲比Hugging Face tokenizers crate低5倍

Perplexity AI 的研究團隊用Rust從頭實現了Unigram分詞器，並將其開源在pplx-garden（推理技術倉庫）中。該分詞器針對XLM-RoBERTa模型（使用SentencePiece訓練的25萬token Unigram詞彙表）進行優化。

在生產輸入長度下，新編碼器的p50延遲比Hugging Face tokenizers crate低約5倍，比SentencePiece（C++）低約2倍，比IREE分詞器（C）低約1.5倍，且熱路徑無堆分配。在生產中，Perplexity推理棧的CPU利用率降低了5-6倍，重排序器延遲減少了數十毫秒。

為什麼分詞成為瓶頸

LLM推理成本通常圍繞GPU工作，但小型模型（如嵌入模型、分類器和重排序器）則不同。這些模型比前沿transformer小兩到三個數量級。重排序器每請求需對數百個候選文檔評分，GPU計算常在個位數毫秒內完成，但每次輸入都需先經過CPU端分詞。當批量較大時，分詞成為總延遲的重要部分。

什麼是Unigram分詞

Unigram分詞由Kudo於2018年提出，在SentencePiece中實現。它將分詞視為最大概率路徑問題：每個詞彙token有學習到的對數概率，分詞器選擇token分數總和最高的分割。該算法使用Viterbi算法（1967年的動態規劃技術），其中字節位置為圖層級，詞彙token為邊。內層循環在每個字節位置遍歷詞彙trie（前綴樹）。在16K token輸入上，內層遍歷執行數十萬次trie轉換，是熱路徑。

Hugging Face實現的瓶頸

在514 token（512+BOS/EOS）輸入下，Hugging Face tokenizers crate存在三個昂貴模式：每匹配一次就進行String::from_utf8和AHashMap查找（514 token時分配7,295次）；每個字節步驟在trie節點進行AHashMap指針追蹤（4次依賴加載）；長輸入時L2緩存抖動（分配DP表和輸出緩衝區）。每token分配約2 KB和18次分配，在長輸入時累積溢出L2緩存。

三項優化

**優化1：雙數組trie** 將Hugging Face的HashMap trie替換為雙數組trie（Aoe, 1989）。該結構將整個trie編碼為兩個平面整數數組base和check。子節點查找：next = base[node] + byte，驗證check[next] == node。只需兩次數組讀取、一次整數加法和一次比較，無需哈希和指針追蹤。25萬詞彙trie佔用約9 MB連續內存，每次編碼的熱工作集約100 KB，適合L2緩存。p50延遲從155 µs降至68 µs，比原始參考快4.8倍。

**優化2：位圖和內聯打包** 用每節點位圖（四個64位字，32字節）取代check數組，記錄256個可能字節中哪些有有效子節點。位圖查找編譯為單個64位字的位測試。同時將所有四個字段（位圖、基址、token ID和分數）打包到單個64字節緩存行中。代價是trie大小從9 MB增至約50 MB，但熱工作集仍約100 KB。p50進一步降低4.5%，L2訪問從4.6K降至1.8K。

**優化3：大頁面支持** 50 MB trie在默認4 KB頁面下跨越約12,000個虛擬頁面，導致TLB缺失。Perplexity使用MAP_HUGETLB標誌通過mmap將trie映射到2 MB大頁面，50 MB僅佔25個頁面。根據輸入長度，延遲降低3-12%，最大增益在4,098 token時達12.0%。

最終基準結果

在Intel Xeon Platinum 8488C單線程，10,000次迭代，514 token下：

• Hugging Face tokenizers crate：349 µs，3.60M指令，7,295次分配
• SentencePiece（C++）：128 µs，1.83M指令，1,559次分配
• IREE分詞器（C）：112 µs，2.28M指令，1次分配
• Perplexity最終版本：~63 µs，1.04M指令，0次分配

完整優化序列使每次編碼指令數從3.66M降至1.04M，降低3.5倍。在長輸入上性能優勢更大。最終編碼器輸出與參考完全一致，生產中通過rayon實現跨核心並行。

開源發佈

Perplexity AI 已將重寫的Unigram分詞器開源在pplx-garden，三項針對性優化去除了熱路徑上的浪費工作。開源社區可立即使用或進一步優化。