Perplexity AIがUnigramトークナイザーをオープンソース化、Hugging Face tokenizers crate比5倍の低レイテンシを達成

Perplexity AIの研究チームは、UnigramトークナイザーをRustでスクラッチから再実装し、そのコードを推論技術リポジトリpplx-gardenでオープンソース化しました。このトークナイザーは、SentencePieceでトレーニングされた25万トークンのUnigram語彙を持つXLM-RoBERTaモデルを対象としています。

本番環境の入力長において、新しいエンコーダはHugging Face tokenizers crateと比較してp50レイテンシを約5倍、SentencePiece（C++）と比較して約2倍、IREEトークナイザー（C）と比較して約1.5倍低減し、ホットパスでのヒープ割り当てはゼロです。本番環境では、Perplexityの推論スタックのCPU使用率を5〜6倍削減し、リランカー（再ランク付け器）のレイテンシを数十ミリ秒削減しました。

なぜトークン化がボトルネックになったか

LLM推論コストは通常GPU作業（KVキャッシュ、アテンションカーネル、エキスパートルーティング）の観点で語られますが、小さなモデル（埋め込みモデル、分類器、リランカーなど）は異なります。これらのモデルはフロンティアトランスフォーマーより2〜3桁小さく、リランカーがリクエストごとに数百の候補文書をスコアリングする場合、GPU計算は一桁ミリ秒で完了しますが、各入力は最初にCPU側のトークン化を通過する必要があります。バッチサイズが大きい場合、トークン化は全レイテンシのかなりの部分を占めます。

Unigramトークン化とは

Unigramトークン化はKudo（2018）によって導入され、SentencePieceに実装されています。これは、各語彙トークンが学習された対数確率を持つ、最大確率パス問題としてセグメンテーションを定式化します。トークナイザーは、トークンスコアの合計が最も高いセグメンテーションを選択します。最適なパスを見つけるアルゴリズムはViterbiアルゴリズム（1967年の動的計画法）で、バイト位置がグラフ層、語彙トークンがエッジとなります。内側のループは各バイト位置で語彙トライ（接頭辞木）を走査します。16Kトークン入力では、この内側の走査は数十万回のトライ遷移を実行し、ホットパスとなります。

Hugging Face実装のボトルネック

514トークン（512+BOS/EOS）入力で、Hugging Face tokenizers crateには3つのコストのかかるパターンがありました。

• マッチごとの割り当て：String::from_utf8 + AHashMapルックアップ（514トークンで7,295回の割り当て、16Kで299,171回）
• バイトごとのポインタ追跡：各トライノードでのAHashMap、ステップごとに4回の依存ロード
• 長入力でのL2スラッシング：DPテーブルと出力バッファが呼び出しごとに新たに割り当てられる（L2ミス率：128トークンで8%、16Kで50%）

トークンごとの割り当ては定数（約2KB、18回の割り当て）であり、長入力で累積するとL2キャッシュをオーバーフローします。

3つの最適化

**最適化1：ダブルアレイトライ** Hugging FaceのHashMapトライをダブルアレイトライ（Aoe, 1989）に置き換えました。これはトライ全体を2つのフラットな整数配列（baseとcheck）でエンコードします。子ノードのルックアップは、next = base[node] + byte、check[next] == nodeを検証するだけです。2回の配列読み取り、1回の整数加算、1回の比較で、ハッシュもポインタ追跡もありません。25万語彙のトライは約9MBの連続メモリに収まり、エンコードごとのホットワーキングセットは約100KBでL2キャッシュに収まります。p50レイテンシが155µsから68µsに低下し、元のリファレンスと比較して4.8倍高速化しました。

**最適化2：ビットマップとキャッシュラインパッキング** ダブルアレイトライはステップごとに2回の依存配列ロードが必要です。Perplexityはcheck配列をノードごとのビットマップに置き換えました。これにより、256の可能なバイト値のうちどの子ノードが有効かを4つの64ビットワード（32バイト）で記録します。ビットマップルックアップは1つの64ビットワードに対する単一ビットテストにコンパイルされます。さらに、4つのフィールド（ビットマップ、ベース、トークンID、スコア）を単一の64バイトキャッシュラインにパックしました。これにより、トライサイズは約9MBから約50MBに増加しますが、ホットワーキングセットは約100KBのままです。p50がさらに4.5%削減され、L2アクセスが4.6Kから1.8Kに減少しました。

**最適化3：ヒュージページの使用** 50MBのトライは、デフォルトの4KBページを使用すると約12,000の仮想ページにまたがり、TLBミスが発生します。PerplexityはMAP_HUGETLBフラグを指定したmmapを使用してトライを2MBのヒュージページにバックアップしました。これにより、同じ50MBが25ページに収まり、TLBに収まります。入力長に応じてレイテンシが3〜12%削減され、最大のゲインは4,098トークンで12.0%でした。

最終ベンチマーク結果

Intel Xeon Platinum 8488C、シングルスレッド、10,000イテレーション、514トークンでの結果：

• 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倍削減されました。短い入力ではウォールクロックも同様の比率で改善し、長い入力では参照実装のトークンごとの割り当てがL2およびL3キャッシュをオーバーフローするため、さらに差が広がります。

最終的なPerplexityエンコーダは参照実装とトークン単位で完全に一致する出力を生成します。本番環境では、rayonを使用してコア間で並列化されています。

オープンソースリリース

Perplexity AIは、UnigramトークナイザーをRustで再実装し、pplx-gardenでオープンソース化しました。3つのターゲット最適化により、ホットパスから無駄な作業を排除しました。コミュニティはすぐに利用したり、さらに最適化を加えることができます。