UnigramトークナイザーのCPUパフォーマンス向上

当社では、パフォーマンス向上を目的としたプロジェクトの一環として、Unigramトークナイザーをゼロから再実装しました。実稼働環境の入力長において、新しいエンコーダーは、Hugging Faceのtokenizersクレートと比較してp50レイテンシを約5倍、SentencePiece (C++) と比較して約2倍、IREEのトークナイザー (C) と比較して約1.5倍削減し、定常状態でのヒープ割り当てをゼロにしました。本番環境では、これにより推論スタックのCPU使用率が5〜6倍削減され、リランカーのレイテンシを数十ミリ秒短縮することに成功しました。

LLMの推論コストは、通常、KVキャッシュ、アテンションカーネル、エキスパートルーティングといったGPUワークとして捉えられています。しかし、実稼働トラフィックの大部分は、フロンティアトランスフォーマーよりも2〜3桁小さい、はるかに小規模なモデルで実行されています。埋め込みモデル、分類器、リランカーも依然としてGPU上で実行されますが、フォワードパスが十分に短いため、その周辺のすべての処理が重要になります。

最も明確な例として、1リクエストあたり数百の候補ドキュメントをスコアリングするリランカーが挙げられます。小規模なモデルでは、GPUの計算は通常数ミリ秒で完了しますが、すべての入力は最初にCPU側でトークン化を行い、テキストを語彙IDに変換する必要があります。バッチサイズが大きい場合、これはリクエストの総レイテンシにおいて大きな割合を占める可能性があります。

本作業では、SentencePieceで学習された250KトークンのUnigram語彙であるXLM-RoBERTaを対象とします。ファインチューニングされたRoBERTaファミリーのエンコーダーは、ランキング、検索、類似性タスクにおいて一般的な選択肢です。ここではシングルスレッドでのエンコードごとのレイテンシを報告します。本番環境ではrayonを使用してコア間で並列化を行っています。

本投稿全体を通じて、Hugging Faceのtokenizersクレートをベンチマークの基準として使用します。これは、多くのチームがデフォルトで使用するRustのトークナイザーであるためです。比較のため、他の2つの主要な実稼働実装であるSentencePiece (C++) とIREE (C) も含めています。

2018年にKudoによって導入され(Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates)、SentencePieceで採用されたUnigramトークン化は、セグメンテーションを「最も確率の高いパス」問題として構成します。各語彙トークンには学習済みの対数確率があり、エンコーダーはトークンスコアの合計が最大となるセグメンテーションを選択します。この問題はビタビアルゴリズム(1967)を用いて解決され、バイト位置をグラフの層、語彙トークンを連続するバイト範囲にまたがるエッジとして扱います。この漸化式は、バイト位置に対する1次元の動的計画法 (DP) です。

python

best = [-inf] * (n + 1)
best[0] = 0
back = [None] * (n + 1)

for i in 0..n:
    if best[i] == -inf:
        continue
    if byte i is a UTF-8 continuation byte:
        continue

    walk the vocabulary trie from byte position i:
        for each matching token t ending at position j:
            cand = best[i] + score(t)
            if cand > best[j]:
                best[j] = cand
                back[j] = (i, token_id(t))

    if no token matched, fall back to <unk>

backtrack from position n using back[] to recover token IDs

外側のループは入力長に対して線形であり、各位置において最長の語彙トークン(XLM-RoBERTaの場合は48 UTF-8バイト)によって制限されます。内側のトライ探索がホットパスであり、16kトークンの入力では、1エンコードあたり数十万回の遷移と数千万回の命令実行が行われます。

Hugging Faceのtokenizers v0.22では、この漸化式が直接実装されており、逆引き用のハッシュマップと、すべてのトライノードにハッシュマップの子ノードが配置されています。

rust

vocab:        Vec<(String, f64)>           // indexed by token ID
token_to_ids: AHashMap<String, u32>        // reverse lookup
trie:         Trie<u8>                     // byte-level prefix trie

struct Node<Label> {
    is_leaf: bool,
    children: AHashMap<Label, Node<Label>>, // heap-allocated map per node
}

そして内側のループ(簡略化)では、接頭辞探索の一致ごとに新しい文字列が割り当てられ、ハッシュマップを探索します。

rust

fn encode_optimized(&self, sentence: &str) -> Result<Vec<String>> {
    // alloc
    let mut best = vec![BestPathNode::default(); sentence.len() + 1];

    let mut starts_at = 0;
    while starts_at < sentence.len() {
        for tok_bytes in self.trie.common_prefix_search(sentence.bytes().skip(starts_at)) {
            // alloc per match
            let token: String = String::from_utf8(tok_bytes).unwrap();
            // hash lookup
            let id = self.token_to_ids.get(&token).unwrap();
            let score = self.vocab[*id as usize].1;
            // DP update at best[starts_at + tok_bytes.len()]
            // ...
        }
        starts_at += mblen;
    }
    // backtrace into Vec<String>, one allocation per output token
}

このコードは読みやすいRustで書かれていますが、ソースレベルでは見えない多くの処理が行われています。次のセクションでは、どこに時間が費やされているかを分解します。

ビタビアルゴリズムが実行される前に、入力はUnicode正規化(NFKC)と単語境界のマーク付けを通過します。プロファイリングによると、我々のシナリオでは前処理は全エンコード時間の5%未満であるため、本投稿の残りの部分ではビタビのフォワードパスに焦点を当てます。

CPUコードを高速化する方法は、命令数を減らすか、IPC (1クロックあたりの命令実行数) を高めるかの2つです。Cachegrindは前者を、perf statは後者を測定します。本投稿のあらゆる最適化は、このいずれかに分類されます。

CPUは、キャッシュミス(L1dは約4サイクル、L3ミスからDRAMへのアクセスは数百サイクル)、アウト・オブ・オーダー実行エンジンが並列化できない依存ロードチェーン、分岐予測ミス(誤予測ごとに約15〜20サイクルのフラッシュ)、およびアロケーター呼び出しの4つの要因で停止します。参照用のUnigramエンジン実装は、これらすべてに該当します。

ボトルネック	メカニズム	測定された影響
一致ごとの割り当て	一致したバイトに対するString::from_utf8と、トライ一致ごとのAHashMap探索。トークンIDはリーフノードに存在しますが、参照実装ではそこに格納されていません。	512トークンで7,295回、1kトークンで14,542回、16kトークンで299,171回の割り当て
バイトごとのポインタチェイス	各トライノードのAHashMapで次のバイトをハッシュ化し、ステップごとに2つのポインタをデリファレンス。Node構造体はis_leafのようなコールドフィールドもバンドルしており、アクセスするたびにキャッシュに引き込まれます。	アロケーターのホットページがL1に存在するため、L1dミス率は低く保たれます(1kトークンで約2.5%)。コストは依存ロードレイテンシとして現れます。
長い入力におけるL2スラッシング	DPテーブル、プレフィックススクラッチ、出力トークンリストが、呼び出しのたびに新しく割り当てられます。	累積割り当てフットプリントがコアあたりの2MB L2キャッシュを超えるため、L2ミス率は128トークンの12%から16kトークンで50%まで上昇します。

参考用トークナイザーのボトルネック、メカニズム、測定された影響。

支配的なコストは、一致ごとの割り当てです。1kトークンの場合、参照実装は1エンコードあたり約7Mの命令を完了させますが、IPCは2.8であり、その大部分はアロケーター呼び出しと文字列変換によるものです。L2ミスは、累積割り当てフットプリントがコアごとのキャッシュから溢れる約1kトークン以降にのみ重要になります。

入力サイズ全体にわたる、参考用トークナイザーのエンコード1回あたりの割り当てバイト数。

トークンあたりの割り当ては、入力サイズに関係なく、約2KBかつ約18回で一定です。レイテンシが非線形になるのはL2キャッシュの影響です。累積バイト数が1kトークン付近で2MBを超えると、ミス率が8%から約50%まで上昇し、p99曲線もそれに合わせて悪化します。

同じUnigramアルゴリズムを実行し、同じXLM-RoBERTaモデルを使用した3つの実稼働トークナイザーを比較します。

エンジン	言語	トライ
Hugging Face	Rust	HashMap子ノードの接頭辞トライ
SentencePiece	C++	Dartsダブル配列トライ
IREEトークナイザー	C	Dartsダブル配列トライ

比較に使用したトークナイザー実装、言語、トライ表現。

SentencePieceとIREEは、Rustラッパーではなく直接C/C++ FFI経由でリンクしています。ラッパーは、基盤となるライブラリのコストには含まれない、呼び出しごとの割り当てや文字列変換を追加してしまうためです。

すべてのエンジンは、完全な最適化(O3、ネイティブターゲット、LTO)でビルドされています。測定は、Intel Xeon Platinum 8488Cの静かな1コアに固定して実行しました。各エンジンと入力サイズに対して、ウォールクロックレイテンシ(p50およびp99、1,000回のウォームアップ後の10,000回反復)、Valgrind Cachegrind経由の退役命令数、perf statハードウェアカウンターからのIPCおよびキャッシュミス率、測定対象ウィンドウのグローバルアロケーターラッパーによるヒープ割り当て回数を収集しました。

514トークンの場合(512にBOS/EOS注入の+2):

実装	p50レイテンシ	命令数	IPC	割り当て回数
Hugging Face	349 µs	3.60M	2.89	7,295
SentencePiece	128 µs	1.83M	3.39	1,559
IREE	112 µs	2.28M	3.06	1

514トークンにおけるトークナイザーのレイテンシ、命令数、IPC、アロケーション。

SPとIREEは、参照実装のレイテンシの約3分の1に収まります。命令数は約半分であり、より高いIPCで実行され、割り当て回数は1〜2桁少なくなります。

SentencePieceとIREEはどちらも、1989年にAoeによって導入されたダブル配列トライを使用しています(Aoe, 1989; Karoonboonyanan)。トライは、baseとcheckという2つのフラットな整数配列にエンコードされています。

cpp

int base[N];   // base[node] = offset for this node's children
int check[N];  // check[child] == node validates the parent-child link

子ノードの検索は次のとおりです。

cpp

int next = base[node] + byte;
if (check[next] == node) { /* valid transition; child node = next */ }

トライ走査の各ステップでは、ハッシュ化、ポインタチェイス、ノードごとのヒープ割り当てなしに、2回の配列ロード、1回の整数加算、1回の比較のみが必要です。構築プロセスでは、各ノードの子供が検証配列の空きスロットに収まるように配置され、続くバイトが異なる限り、複数の親が領域を共有できるようにします。

XLM-RoBERTaの250K語彙では、トライ全体が約9MBの連続メモリに収まります(オフセット、検証、スロットごとのトークンID用に約780Kスロット×12バイト)。単一エンコードのホットワーキングセットは100KBオーダーであり、L2キャッシュに十分に収まり、ハードウェアプリフェッチャーが規則的なアクセスパターンを認識します。

一方、参照実装ではノードごとに子ノードのハッシュマップを保持しており、バイトステップごとに約4回の依存ロードが発生します。バイトのハッシュ化、バケットのロード、バケットエントリへのポインタ追跡、ヒープ上の子ノードへの別のポインタ追跡、そして次のバイトのためのハッシュマップの確認が行われます。ヒープオブジェクトは分散し、ロードの多くはキャッシュミスを引き起こします。

SPとIREEはどちらも、トークンIDをトライノードに直接格納します。参照実装のトライは一致したバイトのみを返すため、エンコーダーは一致ごとに新しい文字列を割り当て、サイドハッシュマップでトークンIDを回復し、別のベクターからスコアを読み取る必要があります。SPはトライに格納された値をIDとして返し、浮動小数点配列からスコアを読み取ります。

cpp

// unigram_model.cc:567
const size_t num_nodes = trie_->commonPrefixSearch(
    begin, trie_results.data(), trie_results.size(),
    static_cast<int>(end - begin));

for (size_t k = 0; k < num_nodes; ++k) {
    const int   id    = trie_results[k].value;   // already in the trie
    const float score = GetScoreInlined(id);     // flat float[] indexed by id
    Lattice::Node *node = lattice->Insert(begin_pos, length);
    node->id = id;
    node->score = score;
}

IREEは1回の呼び出しですべての一致を返す代わりにステートフルなカーソルAPIを使用しますが、コストモデルは同じです。トークンIDとスコアはフラット配列のインデックスとして取得可能であり、トライ探索自体に割り当てやハッシュ化は含まれません。IREEは1エンコードあたり1つの出力バッファのみを割り当てます。一方、SPはDP中に一致ごとにラティスノードを追加で割り当て、514トークンで合計約1,559回の割り当てを行います。

トライを変更する前に、アルゴリズムのコストと実装コストを切り分けます。不要な作業を取り除いた場合、参照実装はどれほど速くなるのでしょうか？

このベースラインでは、参照実装のHashMapトライを維持しつつ、2つの変更を加えます。呼び出し元が所有するスクラッチ構造体にDPテーブル、前処理済みバイト、出力トークンIDを保持することで、ウォームアップ後はホットパスでの割り当てが発生しないようにしました。トークンIDとスコアはトライノードに直接格納され、一致ごとの文字列変換とサイドハッシュマップの探索を排除しています。

内側のループ:

rust

for pos in 0..n {
    if scratch.prep[pos] & 0xC0 == 0x80 { continue; }
    let s0 = scratch.best_score[pos];
    if s0 == f32::NEG_INFINITY { continue; }

    let mut node = 0u32;
    for end in pos..n {
        let byte = scratch.prep[end];
        match self.nodes[node as usize].children.get(&byte) {
            Some(&child) => {
                node = child;
                let n = &self.nodes[node as usize];
                if n.is_terminal {
                    let total = s0 + n.score;
                    if total > scratch.best_score[end + 1] {
                        scratch.best_score[end + 1] = total;
                        scratch.best_start[end + 1] = pos as u32;
                        scratch.best_id[end + 1] = n.token_id;
                    }
                }
            }
            None => break,
        }
    }
}

これは同じトライ形状と同じアクセスパターンを持つ同じアルゴリズムであり、割り当てと冗長な探索がないだけです。514トークンでの結果は以下の通りです。

エンジン	p50	命令数	IPC	割り当て回数
Hugging Face	326 µs	3.66M	2.86	7,295
SentencePiece	127 µs	1.81M	3.55	1,559
IREE	111 µs	1.58M	3.26	1
ベースライン	156 µs	1.53M	2.62	0

514トークンにおける既存のトークナイザー実装と比較した、ゼロアロケーションのベースライン。

一致ごとの文字列具体化と二次ハッシュマップ探索を削除することで、参照実装と比較して命令数を2.4倍削減しました。以前はアロケーター作業の影に隠れていたハッシュマップのポインタチェイスが、バイトあたりの支配的なコストとなったため、それと同時にIPCが低下しています(2.86→2.62)。

入力サイズ別の参考用およびゼロアロケーションのベースライントークナイザーのp50およびp99レイテンシ。

入力サイズ別の参考用およびゼロアロケーションのベースライントークナイザーのL2キャッシュミス率。

ベースラインは、すべてのサイズにおいて参照実装よりもp50レイテンシを2〜3倍削減し、L2ミス率は参照実装の8〜48%に対して2kトークンまで2〜3%近くに保たれています。残るコストはHashMapのポインタチェイスであり、次のセクションで対応します。

HashMapトライはポインタチェイスです。SPとIREEはこれをダブル配列トライで解決しており(前セクション参照)、我々も同じデータ構造を採用します。違いは、SPとIREEはストリーミング入力、出力オフセット、複数のエンコードモードをサポートする汎用ライブラリである点です。その柔軟性のためにトライ周辺の余分な簿記(SPではラティス、IREEでは8段のパイプライン)にコストがかかっています。我々にはそのどれも不要なため、トライをビタビループ内に直接インライン化し、残りを排除しました。

rust

fn child(&self, node: u32, byte: u8) -> Option<u32> {
    let base = self.base[node as usize];
    let next = base + byte as i32;
    if self.check[next as usize] == node as i32 {
        Some(next as u32)
    } else {
        None
    }
}

バイトステップは2回の配列読み込み、整数加算、比較というコストです。トライ全体は約9MBですが、単一エンコードではそのうち約100KBしか走査しないため、ホットデータはL2に留まります。

エンジン	p50	p99	命令数	L2ミス%	L3ミス%
Hugging Face	327 µs	3,340 µs	3.66M	8.17%	14.66%
SentencePiece	127 µs	1,140 µs	1.82M	4.74%	4.07%
IREE	111 µs	132 µs	1.61M	5.82%	3.87%
ベースライン	155 µs	1,170 µs	1.53M	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	6.45%	11.42%

ダブル配列トライ採用後のトークナイザーのレイテンシ、命令数、キャッシュミス率。

ウォールクロックはベースラインから2.3倍、参照実装からは4.8倍低下しました。命令数はベースラインからさらに28%減少しました。ラティスとパイプラインをカットすることで、同じトライ形状においてSPの1.82MおよびIREEの1.61Mに対して、1エンコードあたり1.10Mの命令実行で済むようになりました。その差は、汎用APIのために保持されている簿記処理によるものです。

最適化1のDartsトライには、バイトステップごとに2回の依存ロードが必要です。まず親のbaseオフセット、次に候補子スロットのcheckエントリを確認して親がそれを所有しているかを確認します。checkエントリは所有権を検証するためだけに存在します。同じ疑問は、どのバイトが子供を持つかを記録するノードごとのビットマップでも回答可能です。ダブル配列トライの構成上、親のビットマップにおけるバイトbのビットは、そのバイトの候補スロットが親によって所有されているときに正確にセットされます。ビットマップとcheck配列は同じ情報を持っているため、ランタイム構造からcheck配列を完全に排除できます。

ビットマップ探索は、64ビットワードに対する1回のビットテストにコンパイルされ、約1サイクルで実行されます。レジスタ内のすべてのセットビットをカウントする関連するPOPCNT命令は、他のビットマップインデックス型トライ設計で子オフセットをビットマップから直接計算するために使用されるものです。我々のDarts構成では、そのオフセットがすでにノードのbaseフィールドに格納されているため、これは不要です。

rust

#[repr(C, align(64))]
struct NodeHot {
    bitmap:   [u64; 4],     // 32 B  — "which bytes have children"
    base:     i32,          //  4 B  — child slot offset
    token_id: i32,          //  4 B  — vocab id (-1 if not terminal)
    score:    f64,          //  8 B  — Viterbi log-prob
    _pad:     [u8; 16],     // 16 B  — pad to one cache line
}

// Lookup: one cache-line load, one bit test, one add.
fn child(&self, parent: u32, byte: u8) -> Option<u32> {
    let h = &self.node_hot[parent as usize];
    if !h.bitmap.test(byte) { return None; }
    Some(h.base.wrapping_add(byte as i32) as u32)
}

ついでに、4つのノードごとのフィールド(ビットマップ、base、トークンID、スコア)を1つの64バイトキャッシュラインに詰め込みます。ノードへのアクセスごとにすべてが読み込まれるようになり、ターミナルノードで必要なスコアとトークンIDに加え、次のステップに必要なビットマップとbaseも取得されます。check配列は構築時にのみ使用され、ランタイムレイアウトからは削除されます。ホットパスのインデックスは、トライの構築とループ構造によって境界が保証されているためチェックされません。トレードオフはスペースです。ノードあたり64バイト×約780Kスロットで、トライは9MBから50MBへと増大します。

514トークンでの結果:

エンジン	p50	p99	命令数	L2アクセス数	L2ミス%	L3ミス%
Hugging Face	327 µs	3,340 µs	3.66M	15.0K	8.17%	14.66%
SentencePiece	127 µs	1,140 µs	1.82M	7.5K	4.74%	4.07%
IREE	111 µs	132 µs	1.61M	4.7K	5.82%	3.87%
ベースライン	155 µs	1,170 µs	1.53M	6.1K	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	4.6K	6.45%	11.42%
+ ビットマップとインラインパッキング	65 µs	77 µs	1.06M	1.8K	20.53%	15.01%

ビットマップとインラインパッキングの追加後、514トークンでのトークナイザー結果。

ウォールクロックは514トークンでさらに4.5%低下し、入力長に応じて伸び、4kおよび16kでは6〜7%に達します。check配列を削除したことで、バイトあたりのトライステップが2回のキャッシュラインロードから1回になり、エンコードあたりのL2アクセス数が2.5倍減少しました(標準のDartsの4.6Kに対して1.8K)。ミス率は6%から20%に上昇しますが、その大部分は分母の影響です。アクセス数が半分以下になったため、わずかなミスが総時間に占める割合が大きくなったのです。

トライデータは約50MB(64バイト×780,000ノード)です。デフォルトのLinuxプロセスでは、このメモリは4KBページでバックアップされるため、トライ全体は約12,000の異なる仮想ページに広がり、CPUの第1レベルデータTLBが保持できる量(Sapphire Rapidsは96エントリ)を大幅に超えています。各ビタビステップは異なるノードに触れるため、512トークンのエンコード全体で、エンコードごとの予算の約3%にあたる約9,000サイクルをページテーブル探索に費やすことになります。

修正策は、トライを4KBページではなく2MBのHuge Pagesでバックアップすることです。これにより、同じ50MBが25ページに収まり、TLBに完全に収まるようになり、ページテーブルディレクトリ自体がL2キャッシュを取り合うこともなくなります。

Linuxにはメモリを2MBページでバックアップする方法が2つあります。日和見的なアプローチであるTransparent Huge Pages (THP) では、プログラムがカーネルにメモリ領域を4KBから2MBページにアップグレードするよう要求し、バックグラウンドのカーネルスレッドが時間的・物理的に余裕があるときに実行します。もう一つの選択肢は、カーネルが起動時に確保するあらかじめ予約されたHuge Pageプールから直接トライを割り当てることです。

rust

let raw = unsafe {
    libc::mmap(
        std::ptr::null_mut(),
        alloc_len,
        libc::PROT_READ | libc::PROT_WRITE,
        libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_HUGETLB,
        -1,
        0,
    )
};

このプールはホスト上に存在する必要があります(起動時にvm.nr_hugepagesを設定)。当社の実稼働インスタンスでは10,561個のHuge Pageが予約されており、トライはそのうち24個を消費します。mmapの後、/proc/<pid>/smapsで領域が確認できます: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB。

得られるゲインは入力サイズ全体で3〜12%のウォールクロック短縮であり、ワーキングセット長に比例して安定しています。

トークン	Huge Pagesなし	Huge Pagesあり	Δ
130	15.7 µs	15.2 µs	−3.6%
514	65.4 µs	63.1 µs	−3.4%
1,026	129.5 µs	125.3 µs	−3.3%
4,098	773 µs	679 µs	−12.0%
16,386	3,808 µs	3,554 µs	−6.8%

各入力サイズにおける、巨大ページの有無によるトークナイザーのレイテンシ。

最大の利益は4kトークンで見られます。ワーキングセットがページテーブルトラフィックとトライデータがL2帯域を取り合うほど十分に大きい場合です。4kを超えるとL3ミスが支配的になり、総時間に対するページテーブル探索の割合が減少するため、利得は小さくなります。

これら3つの最適化を統合した最終的なエンジンを、テストした他のすべてのRoBERTaトークナイザーと比較しました。ネイティブのC++ SentencePieceおよびC IREEバイナリに加え、実稼働環境でユーザーが使用する既製のRustラッパーであるsentencepieceおよびiree-tokenizerもベンチマークしました。sentencepieceクレートは、あらゆる入力サイズでネイティブのC++バイナリより1.6〜1.9倍高いレイテンシを示しました。そのオーバーヘッドは、ラッパーがエンコードのたびに新しいトークンピースリストを割り当て、それぞれがピーステキストのコピーを保持していることに起因します。iree-tokenizerクレートは内部バッファを入力サイズに合わせてサイズ変更するため、ネイティブIREEバイナリが入力サイズに関係なく支払う定数的なエンコードごとの割り当てを回避できます。これは割り当てがトークン化作業を支配してしまう短い入力において重要です。実稼働の入力長では定数コストが償却され、Rustにトークンをコピーし戻すラッパーの呼び出しごとのオーバーヘッドが、ネイティブバイナリのレイテンシを上回る原因となります。

トークナイザー実装および入力サイズ全体にわたる、エンコード1回あたりのp50およびp99レイテンシ。

最終的なエンジンは、同じ250,000トークンの語彙で同じビタビアルゴリズムを実行し、ホットパスでは何も割り当てず、参照実装とトークン単位で完全に一致する出力を生成します。各バイトステップは現在、1つの64バイトキャッシュラインロード、1回のビットマップテスト、1回の整数加算であり、ベースライン実装のHashMapプローブ、アロケーター呼び出し、およびいくつかの特殊ケース分岐から削減されました。エンコードあたりの命令数は3.66Mから1.04Mへ、3.5倍減少しました。ウォールクロックも短い入力ではその比率に一致し、参照実装のトークンごとの割り当てがL2およびL3から溢れ始める長い入力では、その差はさらに広がります。

本プロジェクトの成果により、推論スタックのプリプロセス層におけるCPU使用率が5〜6倍削減され、リランカーのレイテンシが数十ミリ秒短縮されました。

レイテンシ削減の大部分は、無駄な作業を排除する軽量な変更によるものでした。エンコードごとの割り当てを排除し、HashMapトライをダブル配列に置き換え、汎用APIのオーバーヘッドを削除したことです。既存のノードごとのビットマップを有効性チェックに使用し、トライを2MBのHuge Pagesでバックアップすることは、入力サイズにもよりますが、ダブル配列トライへの切り替えという大きな変更に加え、10〜20%の改善をもたらしました。

SentencePieceおよびIREEを囲む既製のRustラッパークレートは、同じことを逆の立場で示しています。ネイティブのC/C++トークナイザーの上にラッパーが導入するFFI境界と呼び出しごとの割り当ては、レイテンシの数値に見事に現れています。最終的なビタビエンコーダーのRustソースはgithub.com/perplexityai/pplx-gardenで公開されています。

Hugging Face tokenizers — https://github.com/huggingface/tokenizers

SentencePiece — https://github.com/google/sentencepiece

IREE tokenizer — https://github.com/iree-org/iree

Kudo, T. (2018). Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates. https://arxiv.org/abs/1804.10959

Viterbi, A. J. (1967). Error bounds for convolutional codes and an asymptotically optimum decoding algorithm. IEEE Transactions on Information Theory. https://en.wikipedia.org/wiki/Viterbi_algorithm

Aoe, J. (1989). An efficient digital search algorithm by using a double-array structure. IEEE Transactions on Software Engineering. https://ieeexplore.ieee.org/document/31365/

Karoonboonyanan, T. An Implementation of Double-Array Trie. https://linux.thai.net/~thep/datrie/datrie.html

POPCNT (x86) — https://www.felixcloutier.com/x86/popcnt

Huge pages — https://docs.kernel.org/admin-guide/mm/hugetlbpage.html

BibTeX

@online{perplexity2026ImprovingUnigramTokenizerCpuPerformance,
  author = {{Perplexity} and Vladimir Zaytsev and Bex Brooker and Daisy Lee and Roshan Shankar},
  title = {{Improving unigram tokenizer CPU performance}},
  date = {2026-05-20},
  url = {https://www.perplexity.ai/hub/blog/improving-unigram-tokenizer-cpu-performance}
}

動機

ビタビアルゴリズム

参照実装

前処理に関する注記

参照実装が時間を費やしている箇所

測定方法

SentencePieceとIREEが高速である理由: ダブル配列トライ

ベースライン: 参照実装のゼロ割り当てポート

最適化1: 余分な抽象化を排除したダブル配列トライ

最適化 2: ビットマップとインラインパッキング

最適化 3: トライ木のための巨大ページ（Huge Pages）

最終比較

結びの言葉

参考文献

引用