Zlepšení výkonu procesoru u unigramového tokenizátoru

Náš Unigram tokenizátor jsme kompletně přepracovali od základu v rámci projektu zaměřeného na výkon. Při produkčních délkách vstupu nový kodér snižuje latenci p50 přibližně 5× oproti knihovně tokenizerů od Hugging Face, cca 2× oproti SentencePiece (C++) a cca 1,5× oproti tokenizeru IREE (C), přičemž nevyužívá žádné alokace haldy ve stabilním stavu. V produkci to snížilo využití CPU v naší inferenční vrstvě 5–6krát a zkrátilo latenci rerankeru o desítky milisekund.

Náklady na inferenci LLM se obvykle popisují jako práce GPU: KV mezipaměti, pozornostní kernely, routování expertů. Nicméně významná část produkčního provozu běží na mnohem menších motelech, o dva až tři řády menších než špičkové transformery. Vrstvy pro embedding, klasifikátory a rerankery stále běží na GPU, ale dopředný průchod je dostatečně krátký na to, aby začalo záležet na všem kolem.

Nejjasnějším příkladem je reranker, který ohodnocuje stovky kandidátních dokumentů na jeden požadavek. U malého modelu se výpočet na GPU obvykle dokončí v řádu jednotek milisekund, ale každý vstup musí nejprve projít tokenizací na straně CPU, která převádí text na ID slovníku. Když je dávka velká, může to představovat významný podíl celkové latence požadavku.

Tato práce se zaměřuje na model XLM-RoBERTa, který využívá 250ktokenový slovník Unigram natrénovaný pomocí SentencePiece. Jemně doladěné kodéry rodiny RoBERTa jsou běžnou volbou pro úlohy ranking, vyhledávání a určování podobnosti. V celém textu uvádíme latenci pro jednovláknové kódování; v produkci provádíme paralelizaci napříč jádry pomocí rayon.

V tomto příspěvku budeme považovat tokenizers crate od Hugging Face za referenci, vůči níž provádíme srovnání, protože jde o výchozí rustový tokenizátor, po kterém většina týmů sáhne. SentencePiece (C++) a IREE (C) jsou zahrnuty pro srovnání jako další dvě hlavní produkční implementace.

Tokenizace Unigram, kterou představil Kudo v roce 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) a která je obsažena v SentencePiece, chápe segmentaci jako problém hledání nejpravděpodobnější cesty: každý token slovníku má naučenou log-pravděpodobnost a kodér vybírá takovou segmentaci, jejíž skóre tokenů dává v součtu nejvyšší hodnotu. Problém je řešen pomocí Viterbiho algoritmu (1967), přičemž pozice bajtů představují vrstvy grafu a tokeny slovníku tvoří hrany pokrývající souvislý rozsah bajtů. Rekurze je jednorozměrné dynamické programování (DP) nad pozicemi bajtů:

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

Vnější smyčka je lineární vzhledem k délce vstupu, na každé pozici ohraničená nejdelším tokenem slovníku (48 UTF-8 bajtů pro XLM-RoBERTa). Vnitřní průchod trie strukturou (trie walk) je kritickou cestou: na vstupu o délce 16 tisíc tokenů provede stovky tisíc přechodů a zpracuje desítky milionů instrukcí na jedno zakódování.

Knihovna tokenizers v0.22 od Hugging Face implementuje rekurzi přímo, a to pomocí hashovací tabulky pro zpětné vyhledávání a vnořených hashovacích tabulek dětí v každém uzlu trie:

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
}

A vnitřní smyčka (zjednodušená) – každá shoda při vyhledávání prefixu alokuje nový řetězec a poté dotazuje hashovací tabulku:

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
}

Kód je napsán v čitelném jazyce Rust, ale na úrovni zdrojového kódu je neviditelná spousta probíhající práce. Následující sekce rozebírá, kam všude čas utíká.

Než se spustí Viterbiho algoritmus, vstup projde normalizací Unicode (NFKC) a označením hranic slov. Profilování ukazuje, že předzpracování tvoří méně než 5 % celkového času kódování v našich scénářích, takže se ve zbytku příspěvku zaměříme na dopředný průchod Viterbiho algoritmu.

Existují dva způsoby, jak zrychlit kód pro CPU: provádět méně instrukcí nebo dosáhnout vyššího IPC. Nástroj Cachegrind měří to první, perf stat měří to druhé. Každá optimalizace v tomto příspěvku přispívá k jednomu nebo druhému.

Procesorové jednotky (CPU) se zasekávají ze čtyř důvodů: minice v cache (L1d ~4 cykly, minulá L3 cache s přístupem do DRAM stovky), řetězy závislých načtení, které engine s mimořádným pořadím (out-of-order) nedokáže paralelizovat, nesprávné predikce větví (~15–20 cyklů zahozených na každé špatné odhadnutí) a volání alokátoru. Referenční implementace unigramového enginu naráží na všechny čtyři:

Úzké hrdlo	Mechanismus	Naměřený dopad
Alokace na jednu shodu	String::from_utf8 nad shodnými bajty plus AHashMap vyhledávání pro každou shodu v trie; ID tokenu je již obsaženo v listu, ale reference ho tam neukládá	7 295 alokací při 512 t., 14 542 při 1k, 299 171 při 16k
Pronásledování ukazatelů na jeden bajt	AHashMap v každém uzlu trie, která hashuje následující bajt a dereferencuje dva ukazatele na krok; struktura Node navíc sdružuje studená pole, jako je is_leaf, čímž je při každé návštěvě natahuje do cache	Míra minulých L1d cache zůstává nízká (~2,5 % při 1k), protože horké stránky alokátoru žijí v L1; náklady se projevují jako latence závislého načtení
Tříštění L2 cache u dlouhých vstupů	Tabulka DP, pracovní prostor prefixu a výstupní seznam tokenů jsou při každém volání nově alokovány	Míra minulých L2 cache stoupá z 12 % u 128 tokenů na 50 % u 16k, protože kumulativní alokační stopa přesahuje 2MB L2 cache na jádro

Reference tokenizérová úzká místa, mechanismy a naměřené dopady.

Hlavní náklad představuje alokace na jednu shodu. Při 1 tisíci tokenů referenční implementace provede (dokončí) přibližně 7 milionů instrukcí na jedno zakódování při IPC 2,8, přičemž většina z nich připadá na volání alokátoru a převody řetězců. Minulé L2 cache hrají roli až po zhruba 1 tisíci tokenů, kdy kumulativní alokační stopa přeteče cache jednoho jádra.

Latence reference tokenizéru p50 a p99 podle velikosti vstupu.

Míra chyb v mezipaměti L2 u reference tokenizéru podle velikosti vstupu.

Počet bajtů alokovaných na jedno kódování referenčním tokenizérem napříč velikostmi vstupu.

Alokace na jeden token je konstantní: přibližně 2 KB a ~18 alokací na token bez ohledu na velikost vstupu. Nelinearita latence pochází z L2: kumulativní bajty překračují 2 MB kolem 1 tisíce tokenů, míra minulých cache stoupá z 8 % na ~50 % a křivka p99 se tomu přizpůsobuje.

Porovnáváme tři produkční tokenizátory, které všechny používají stejný unigramový algoritmus nad stejným modelem XLM-RoBERTa:

Engine	Jazyk	Trie
Hugging Face	Rust	Prefixový trie s dětmi v HashMap
SentencePiece	C++	Darts double-array trie
IREE tokenizer	C	Darts double-array trie

Implementace tokenizérů, jazyky a reprezentace trie použité v porovnání.

SentencePiece a IREE propojujeme prostřednictvím přímého C/C++ FFI namísto jejich rustových wrapperů; wrappery přidávají alokace na jedno volání a konverze řetězců, které nejsou součástí nákladů podkladové knihovny.

Všechny enginy jsou sestaveny s plnými optimalizacemi (O3, cílení na nativní architekturu, optimalizace v době linkování). Měření probíhala na jednom vyhrazeném klidném jádru procesoru Intel Xeon Platinum 8488C. Pro každý engine a velikost vstupu sbíráme reálnou latenci (p50 a p99, 10 000 iterací po 1 000 zahřívacích), dokončené instrukce prostřednictvím nástroje Valgrind cachegrind, IPC a míry minulých cache z hardwarových čítačů perf stat a alokace haldy počítané obaleným globálním alokátorem v měřeném okně.

Při 514 tokenech (+2 nad 512 je injektáž BOS/EOS):

Implementace	Latence p50	Instrukce	IPC	Alokace
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

Latence tokenizéru, počet instrukcí, IPC a alokace při 514 tokenech.

SP a IREE dosahují přibližně třetiny latence reference. Dokončují zhruba poloviční počet instrukcí, běží s vyšším IPC a alokují o jeden až dva řády méně.

SentencePiece i IREE používají dvojpole trie, které představil Aoe v roce 1989 (Aoe, 1989; Karoonboonyanan). Trie je zakódováno ve dvou plochých celočíselných polích, base a check:

cpp

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

Vyhledávání potomka probíhá takto:

cpp

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

Každý krok průchodu trie vyžaduje pouze dvě načtení z pole, jedno sčítání celých čísel a jedno porovnání, bez hashování, pronásledování ukazatelů a alokace haldy na uzel. Konstrukce umisťuje každý uzel tak, aby jeho potomci padli na volné pozice v ověřovacím poli, což umožňuje vícero rodičům sdílet oblasti, dokud se jejich potomci liší v následujícím bajtu.

Pro slovník modelu XLM-RoBERTa o velikosti 250 tisíc slov se celý trie vejde do ~9 MB souvislé paměti (~780 tisíc slotů × 12 bajtů pro offset, ověření a ID tokenu na slot). Horká pracovní množina jednoho zakódování se pohybuje v řádu 100 KB, pohodlně v L2 cache, a hardwarový přednačítač (prefetcher) se zaměřuje na pravidelný vzor přístupu.

Naproti tomu reference udržuje hashovací tabulku dětí pro každý uzel, což stojí zhruba čtyři závislá načtení na jeden krok bajtu. Zahashuje se bajt, načte se kbelík (bucket), následuje ukazatel na položku v kbelíku, další ukazatel na podřízený uzel na haldě a nakonec se zkontroluje jeho vlastní hashovací tabulka pro další bajt. Objekty na haldě jsou roztroušené a většina načtení míjí cache.

SP a IREE ukládají ID tokenu přímo v uzlu trie. Referenční implementace trie vrací pouze shodné bajty, takže kodér musí alokovat nový řetězec pro každou shodu, vyhledat ji v boční hashovací tabulce pro zjištění ID tokenu a přečíst skóre ze samostatného vektoru. SP vrací ID jako uloženou hodnotu trie a čte skóre z plochého pole s plovoucí řádovou čárkou:

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 používá stavové rozhraní API s kurzorem namísto vracení všech shod v jediném volání, ale model nákladů zůstává stejný: ID tokenů a skóre jsou přístupné jako indexy v plochém poli, bez alokace nebo hashování při samotném průchodu trie. IREE alokuje pouze jeden výstupní buffer na jedno zakódování; SP navíc alokuje uzel mřížky na každou shodu během DP, což celkově představuje přibližně 1 559 alokací při 514 tokenech.

Než změníme trie, izolujeme náklady algoritmu od nákladů implementace: jak rychle běží reference po odstranění zbytečné práce?

Základní linie zachovává referenční trie s dětmi v hashovací tabulce, ale provádí dvě změny. Pracovní struktura (scratch struct) vlastněná volajícím uchovává tabulku DP, předzpracované bajty a výstupní ID tokenů, takže kritická cesta po zahřátí neprovádí žádné alokace. ID tokenu a skóre jsou uloženy přímo v uzlu trie, čímž odpadá převod řetězců na shodu a vyhledávání v boční hashovací tabulce.

Vnitřní smyčka:

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,
        }
    }
}

Jedná se o stejný algoritmus se stejným tvarem trie a stejným vzorem přístupu, pouze bez alokací a nadbytečných vyhledávání. Výsledky při 514 tokenech:

Engine	p50	Instrukce	IPC	Alokace
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
Baseline	156 µs	1,53M	2,62	0

Základní varianta bez alokací ve srovnání se stávajícími implementacemi tokenizérů při 514 tokenech.

Odstranění materializace řetězců na shodu a sekundárního vyhledávání v hashovací tabulce snižuje počet instrukcí 2,4krát oproti referenci. IPC s nimi klesá (2,86 → 2,62), protože pronásledování ukazatelů v hashovací tabulce, dříve skryté za prací alokátoru, je nyní hlavní náklad na jeden bajt.

Latence reference tokenizéru a základní varianty bez alokací p50 a p99 podle velikosti vstupu.

Míra chyb v mezipaměti L2 u reference tokenizéru a základní varianty bez alokací podle velikosti vstupu.

Základní linie snižuje latenci p50 2–3krát oproti referenci při každé velikosti a míra minulých L2 cache zůstává blízko 2–3 % do 2 tisíc tokenů oproti 8–48 % u reference. Zbývající náklady představuje pronásledování ukazatelů v HashMap, což řeší následující sekce.

Trie s hashovací tabulkou je pronásledování ukazatelů. SP a IREE to řeší pomocí dvojpole trie (viz předchozí sekce) a my si vypůjčujeme stejnou datovou strukturu. Rozdíl spočívá v tom, že SP a IREE jsou obecné knihovny podporující streamovaný vstup, výstupní offsety a více režimů kódování; tato flexibilita je stojí dodatečnou režii kolem trie (mřížka v SP, osmifázová pipeline v IREE). My nic z toho nepotřebujeme, takže vkládáme trie přímo do Viterbiho smyčky a zbytek odstraňujeme.

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
    }
}

Jeden krok bajtu stojí dvě čtení z pole, sčítání celých čísel a porovnání. Celý trie má kolem 9 MB, ale jakékoli jedno zakódování prochází pouze ~100 KB z něj, takže horká data zůstávají v L2 cache.

Engine	p50	p99	Instrukce	Minulé L2 %	Minulé 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 %
Baseline	155 µs	1 170 µs	1,53M	2,59 %	16,66 %
+ Darts	68 µs	80 µs	1,10M	6,45 %	11,42 %

Latence tokenizéru, instrukce a míra chyb v mezipaměti po zavedení double-array trie.

Reálný čas klesá 2,3krát oproti základní linii a 4,8krát oproti referenci. Instrukce klesají o dalších 28 % oproti základní linii. Odstranění mřížky a pipeline nám umožňuje provést 1,10 milionu instrukcí na jedno zakódování oproti 1,82M u SP a 1,61M u IREE při stejném tvaru trie; rozdíl představuje režie, kterou udržují pro obecná API, jež nepotřebujeme.

Darts trie z optimalizace 1 vyžaduje dvě závislá načtení na jeden krok bajtu. Nejprve základní offset rodiče a poté kontrolní záznam v pozici kandidátského potomka, aby se potvrdilo, že ho rodič vlastní. Kontrolní záznam ověřuje pouze vlastnictví. Na stejnou otázku lze odpovědět pomocí bitmapy pro každý uzel, která zaznamenává, které bajty mají potomky. Podle konstrukce dvojpole trie je bit pro bajt b v bitmapě rodiče nastaven přesně tehdy, když kandidátský slot pro tento bajt vlastní rodič. Bitmapa a kontrolní pole nesou stejnou informaci, takže můžeme kontrolní pole z běhové struktury zcela odstranit.

Vyhledávání v bitmapě se zkompiluje do jediného testu bitu proti jednomu 64bitovému slovu, což běží zhruba jeden cyklus. Související instrukce POPCNT, která počítá všechny nastavené bity v registru, je to, co jiné návrhy trie indexovaných bitmapou používají k přímému výpočtu offsetu potomka z bitmapy. My ji nepotřebujeme, protože konstrukce Darts již tento offset ukládá v poli base uzlu.

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)
}

Když už v tom jsme, zabalíme čtyři pole na uzel (bitmapu, base, ID tokenu a skóre) do jednoho 64bajtového řádku cache. Jedna návštěva uzlu tak natáhne všechno: skóre a ID tokenu, které potřebujeme v koncových uzlech, plus bitmapu a base pro další krok. Kontrolní pole se používá pouze během konstrukce a je odstraněno z běhové struktury. Indexování na kritické cestě je nekontrolované, protože konstrukce trie a struktura smyčky již zajišťují meze. Kompromisem je Prostor: při 64 bajtech na uzel × ~780 tisíc slotů roste trie z ~9 MB na ~50 MB.

Výsledky při 514 tokenech:

Engine	p50	p99	Instrukce	Přístupy do L2	Minulé L2 %	Minulé 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 %
Baseline	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 %
+ Bitmapa a vnořené balení	65 µs	77 µs	1,06M	1,8K	20,53 %	15,01 %

Výsledky tokenizéru při 514 tokenech po přidání bitmapy a inline balení.

Reálný čas systematicky klesá o další 4,5 % při 514 tokenech a roste s délkou vstupu, přičemž dosahuje 6–7 % při 4k a 16k. Odstranění kontrolního pole mění krok trie na jeden bajt ze dvou načtení řádků cache na jedno, což znamená 2,5krát méně přístupů do L2 na jedno zakódování (1,8K oproti 4,6K u obyčejného Darts). Míra minulých cache vyskočí ze 6 % na 20 %, ale velká část z toho je dána jmenovatelem: s méně než polovičním počtem přístupů zabírají chyby, které máme, větší podíl.

Data trie zabírají přibližně 50 MB (780 000 uzlů po 64 bajtech). V běžném linuxovém procesu je tato paměť podložena 4KB stránkami, takže celý trie pokrývá asi 12 000 různých virtuálních stránek, což zdaleka přesahuje kapacitu vyrovnávací paměti TLB (Translation Lookaside Buffer) první úrovně procesoru (Sapphire Rapids jí má 96 položek). Každý krok Viterbiho algoritmu se dotýká jiného uzlu a během zakódování 512 tokenů propálíme zhruba 9 000 cyklů prohledáváním tabulek stránek, což představuje asi 3 % rozpočtu na jedno zakódování.

Řešením je podložit trie 2MB obřími stránkami namísto 4KB stránek. Stejných 50 MB nyní zabírá 25 stránek, vejde se celé do vyrovnávací paměti TLB a samotné adresáře tabulek stránek přestanou konkurovat trie o L2 cache.

Linux nabízí dva způsoby, jak podložit paměť 2MB stránkami. Oportunistickým způsobem jsou transparentní obří stránky (THP): váš program požádá kernel o povýšení oblasti paměti ze 4KB stránek na 2MB a na pozadí běžící vlákno kernelu se to pokusí provést, když má čas a k dispozici souvislou fyzickou paměť. Druhou alternativou je alokovat trie přímo z předem vyhrazeného fondu obřích stránek, který kernel vyčlení při startu:

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,
    )
};

Fond musí na hostiteli existovat (vm.nr_hugepages nastaveno při startu). V našich produkčních instancích je rezervováno 10 561 obřích stránek; trie jich spotřebuje 24. Po mmap potvrzuje oblast soubor /proc/<pid>/smaps: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Zisk činí 3–12 % reálného času napříč velikostmi vstupů, což je konzistentní a úměrné délce pracovní množiny:

Tokeny	Bez obřích stránek	S obřími stránkami	Δ
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 %

Latence tokenizéru s obrovskými stránkami i bez nich napříč velikostmi vstupu.

Největší přínos je u 4 tisíc tokenů, kde je pracovní množina dostatečně velká na to, aby provoz tabulek stránek aktivně konkuroval datům trie o šířku pásma L2. Nad 4k se zisk zmenšuje, protože dominují minulé L3 cache a průchody tabulkami stránek tvoří menší podíl celkového času.

Když spojíme tři optimalizace dohromady, zde je finální engine naměřený proti každému jinému RoBERTa tokenizátoru, který jsme testovali. Vedle nativních C++ binárek SentencePiece a C binárek IREE jsme otestovali také hotové rustové balíčky (wrapper crates), po kterých by produkční uživatel sáhl: sentencepiece a iree-tokenizer. Balíček sentencepiece vykazuje o 1,6–1,9krát vyšší latenci než nativní C++ binárka při každé velikosti vstupu, přičemž režie pochází z toho, že wrapper alokuje při každém zakódování nový seznam částí tokenu, přičemž každá drží svou vlastní kopii textu. Balíček iree-tokenizer přizpůsobuje velikost svého interního bufferu každému vstupu, což zabraňuje konstantní alokaci na jedno zakódování, kterou nativní binárka IREE platí bez ohledu na velikost vstupu – to je nejdůležitější u krátkých vstupů, kde by alokace jinak dominovala skutečné práci tokenizace. Při produkčních délkách vstupu se tyto konstantní náklady amortizují a režie samotného wrapperu z kopírování tokenů zpět do Rustu v důsledku volání posouvá jeho latenci nad úroveň nativní binárky.

Latence p50 a p99 na jedno kódování napříč implementacemi tokenizérů a velikostmi vstupu.

Finální engine spouští stejný Viterbiho algoritmus nad stejným 250 000tokenovým slovníkem, na kritické cestě nealokuje nic a produkuje výstup shodný s tokeny reference. Každý krok bajtu je nyní jedním 64bajtovým načtením řádku cache, jedním testem bitmapy a jedním sčítáním celých čísel, což je pokles oproti sondování HashMap, volání alokátoru a několika speciálním větvím v základní implementaci. Instrukce na jedno zakódování klesly z 3,66M na 1,04M, což představuje 3,5násobný pokles. Reálný čas odpovídá tomuto poměru u krátkých vstupů a dále se rozšiřuje u dlouhých vstupů, kde alokace na jeden token u reference začínají přetékat L2 a L3 cache.

V praxi tato práce snížila využití CPU v předzpracovávací vrstvě naší inferenční vrstvy 5–6krát a zkrátila latenci rerankeru o desítky milisekund.

Většina snížení latence pocházela z lehkých změn, které odstranily zbytečnou práci: zrušení alokací na jedno zakódování, nahrazení trie s dětmi v HashMap dvojpoletem a odstranění režie obecného API. Použití stávající bitmapy na uzel jako kontroly platnosti a podložení trie 2MB obřími stránkami dohromady tvořilo 10–20 % navíc k většímu přechodu na dvojpole trie, v závislosti na velikosti vstupu.

Hotové rustové wrappery kolem SentencePiece a IREE ilustrují totéž obráceně. Na vrcholu nativních tokenizátorů v C a C++ jsou v číslech latence vidět hranice FFI a alokace na jedno volání, které wrappery přinášejí. Zdrojový kód v Rustu pro finální Viterbiho kodér je k dispozici na adrese 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}
}

Motivace

Viterbiho algorismus

Referenční implementace

Poznámka k předzpracování

Kam referenční implementace investuje svůj čas

Jak jsme měřili

Proč jsou SentencePiece a IREE rychlé: Dvojpole trie (Double-Array Trie)

Základní linie (Baseline): Port reference s nulovými alokacemi

Optimalizace 1: Dvojpole trie bez dodatečných abstrakcí

Optimalizace 2: Bitmapové a řádkové balení

Optimalizace 3: Obří stránky pro trie

Závěrečné srovnání

Závěrečné myšlenky

Odkazy

Citace