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，三项针对性优化去除了热路径上的浪费工作。开源社区可立即使用或进一步优化。