Zlepšenie výkonu procesora pre unigramový tokenizátor

Náš tokenizátor Unigram sme od základu prepracovali ako zameraný projekt na výkon. Pri dĺžkach vstupov v produkcii nový encoder znižuje latenciu p50 približne 5× v porovnaní s balíkom Hugging Face tokenizers, ~2× v porovnaní so SentencePiece (C++) a ~1,5× v porovnaní s tokenizátorom IREE (C), pričom nevykonáva žiadne alokácie haldy v ustálenom stave. V produkcii to znížilo využitie CPU v našom inferenčnom zásobnáku 5- až 6-násobne a skrátilo latenciu pre-rankera o dvojciferné milisekundy.

Náklady na inferenciu LLM sa zvyčajne vnímajú ako práca GPU: vyrovnávaciu pamäť KV, jadrá pozornosti, smerovanie expertov. Značná časť produkčnej prevádzky však beží na oveľa menších modeloch, o dva až tri rády menších ako modely transformátorov typu frontier. Modely na tvorbu vkladov, klasifikátory a pre-rankery stále bežia na GPU, ale dopredný prechod je dostatočne krátky na to, aby všetko okolo neho začalo byť dôležité.

Najjasnejším prípadom je pre-ranker hodnotiaci stovky kandidátskych dokumentov na jednu požiadavku. Pri malom modeli sa výpočet GPU zvyčajne dokončí v ráde jednotiek milisekúnd, ale každý vstup musí najskôr prejsť tokenizáciou na strane CPU, čím sa text premení na ID slovníka. Keď je dávka veľká, môže to predstavovať významný podiel z celkovej latencie požiadavky.

Táto práca sa zameriava na model XLM-RoBERTa, ktorý využíva slovník Unigram s 250 tisíc tokenmi natrénovaný pomocou SentencePiece. Doladené encodery rodiny RoBERTa predstavujú bežnú voľbu pre úlohy hodnotenia, vyhľadávania a podobnosti. V celom článku uvádzame jednovláknovú latenciu na jeden prechod; v produkcii prácu paralelizujeme naprieč jadrami pomocou rayon.

V tomto príspevku považujeme tokenizers crate od Hugging Face za referenciu, voči ktorej porovnávame výkon, pretože ide o predvolený rustový tokenizátor, po ktorom siahla väčšina tímov. SentencePiece (C++) a IREE (C) sú zahrnuté na porovnanie ako ďalšie dve hlavné produkčné implementácie.

Tokenizácia Unigram, ktorú predstavil Kudo v roku 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) a ktorá bola integrovaná do SentencePiece, chápe segmentáciu ako problém cesty s najvyššou pravdepodobnosťou: každý token slovníka má naučenú log-pravdepodobnosť a encoder vyberá takú segmentáciu, ktorej skóre tokenov dáva v súčte najvyššiu hodnotu. Problém sa rieši pomocou Viterbiho algoritmu (1967), pričom pozície bajtov slúžia ako vrstvy grafu a tokeny slovníka ako hrany pokrývajúce súvislý rozsah bajtov. Rekurencia predstavuje jednorozmerné dynamické programovanie nad pozíciami bajtov:

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

Vonkajšia slučka je lineárna vzhľadom na dĺžku vstupu, pričom na každej pozícii je ohraničená najdlhším tokenom slovníka (48 bajtov UTF-8 pre XLM-RoBERTa). Vnútorný prechod cez štruktúru trie je kritickou cestou: na vstupe s 16k tokenmi vykoná stovky tisíc prechodov a spracuje desiatky miliónov inštrukcií na jeden prechod.

Knižnica tokenizers v0.22 od Hugging Face implementuje rekurenciu priamo, a to s hashovacou mapou pre spätné vyhľadávanie a deťmi v hashovacej mape na každom uzle štruktúry 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 vnútorná slučka (zjednodušená) – každá zhoda pri vyhľadávaní predpony alokuje nový reťazec a potom preveruje hashovaciu mapu:

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 napísaný v čitateľnom jazyku Rust, ale na pozadí prebieha množstvo práce, ktorá nie je na úrovni zdrojového kódu viditeľná. Nasledujúca časť rozoberá, kam sa stráca čas.

Pred spustením Viterbiho algoritmu prechádza vstup normalizáciou Unicode (NFKC) a značením hraníc slov. Profilovanie ukazuje, že predbežné spracovanie zaberá v našich scenároch menej ako 5 % celkového času kódovania, takže v zvyšnej časti príspevku sa zameriavame na dopredný prechod Viterbiho algoritmu.

Existujú dva spôsoby, ako zrýchliť kód v CPU: vykonať menej inštrukcií alebo dosiahnuť vyšší ukazovateľ IPC. Cachegrind meria prvý ukazovateľ, zatiaľ čo perf stat meria druhý. Každá optimalizácia v tomto príspevku prispieva k onej alebo druhej metrike.

Procesory spomaľujú štyri faktory: vyrovnávacia pamäť minie (L1d ~4 cykly, minutie L3 vedúce do DRAM stojí stovky), reťazce závislého načítania, ktoré engine s mimo-poradovým vykonávaním nedokáže paralelizovať, nesprávne predikcie vetvenia (~15–20 cyklov vymazaných pri každom nesprávnom odhade) a volania alokátora. Referenčná implementácia unigramového enginu naráža na všetky štyri:

Úzke miesto	Mechanizmus	Nameraný dopad
Alokácia na zhodu	String::from_utf8 pre zhodné bajty plus vyhľadávanie v AHashMap pre každú zhodu v trie; ID tokenu sa už nachádza v liste, ale referencia ho tam neuloží	7 295 alokácií pri 512 tokenoch, 14 542 pri 1k, 299 171 pri 16k
Hľadanie ukazovateľa na jeden bajt	Štruktúra AHashMap na každom uzli trie, ktorá hashuje nasledujúci bajt a dereferencuje dva ukazovatele na krok; štruktúra Node navyše zahŕňa studené polia, ako je is_leaf, čím ich pri každej návšteve ťahá do vyrovnávacej pamäte	Miera minutia L1d zostáva nízka (~2,5 % pri 1k), pretože horúce stránky alokátora sídlia v L1; náklady sa prejavujú ako latencia závislého načítania
Hromadenie chýb v L2 pri dlhých vstupoch	Tabuľka DP, dočasná pamäť pre predpony a zoznam výstupných tokenov sa pri každom volaní nanovo alokujú	Miera minutia L2 stúpa z 12 % pri 128 tokenoch na 50 % pri 16k, keďže kumulatívna stopa alokácií presahuje 2 MB L2 pre každé jadro

Referenčné úzke miesta tokenizátora, mechanizmy a namerané dopady.

Dominantné náklady predstavuje alokácia na zhodu. Pri 1k tokenoch referenčná implementácia dokončí približne 7 miliónov inštrukcií na jeden prechod pri IPC 2,8, pričom väčšina z nich pripadá na volania alokátora a konverzie reťazcov. Minutia L2 začínajú byť dôležité až po ~1k tokenoch, akonáhle kumulatívna stopa alokácií pretečie vyrovnávaciu pamäť jadra.

Referenčný tokenizátor p50 a p99 latencia podľa veľkosti vstupu.

Miera zlyhaní L2 cache referenčného tokenizátora podľa veľkosti vstupu.

Počet bajtov alokovaných na jedno kódovanie referenčným tokenizátorom naprieč veľkosťami vstupu.

Alokácia na jeden token je konštantná: asi 2 KB a približne 18 alokácií na token bez ohľadu na veľkosť vstupu. Nelinearita latencie vyplýva z L2: kumulatívne bajty prekročia 2 MB okolo 1k tokenov, miera minutia stúpne z 8 % na ~50 % a krivka p99 sa tomu prispôsobí.

Porovnávame tri produkčné tokenizátory, pričom všetky spúšťajú rovnaký algoritmický Unigram oproti rovnakému modelu XLM-RoBERTa:

Engine	Jazyk	Trie
Hugging Face	Rust	Prefixová štruktúra trie s deťmi typu HashMap
SentencePiece	C++	Dvojpoľová štruktúra trie Darts
Tokenizátor IREE	C	Dvojpoľová štruktúra trie Darts

Implementácie tokenizátora, jazyky a reprezentácie trie použité v porovnaní.

SentencePiece a IREE prepájame prostredníctvom priameho rozhrania C/C++ FFI namiesto ich rustových obalov; obaly pridávajú alokácie a konverzie reťazcov na jedno volanie, ktoré nie sú súčasťou nákladov základnej knižnice.

Všetky enginy sú zostavené s plnými optimalizáciami (O3, natívne cielenie, optimalizácia pri prepojení). Merania sú spustené a pripnuté na jedno tiché jadro na procesore Intel Xeon Platinum 8488C. Pre každý engine a veľkosť vstupu zbierame latenciu reálneho času (p50 a p99, 10 000 iterácií po 1 000 zahrievacích), dokončené inštrukcie prostredníctvom Valgrind cachegrind, hodnoty IPC a mieru minutia vyrovnávacej pamäte z hardvérových počítadiel perf stat a alokácie haldy počítané obalom globálneho alokátora okolo meraného okna.

Pri 514 tokenoch (+2 nad 512 predstavuje vloženie BOS/EOS):

Implementácia	Latencia p50	Inštrukcie	IPC	Alokácie
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

Latencia tokenizátora, počet inštrukcií, IPC a alokácie pri 514 tokenoch.

SP a IREE dosahujú približne tretinu latencie referencie. Dokončujú približne o polovicu menej inštrukcií, behajú pri vyššom IPC a alokujú o jeden až dva rády menej.

SentencePiece a IREE využívajú dvojpoľovú štruktúru trie, ktorú predstavil Aoe v roku 1989 (Aoe, 1989; Karoonboonyanan). Štruktúra trie je zakódovaná v dvoch plochých celočíselných poliach, 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

Vyhľadanie dieťaťa prebieha takto:

cpp

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

Každý krok prechádzania trie vyžaduje iba dve načítania z poľa, jedno celočíselné sčítanie a jedno porovnanie, bez hashovania, naháňania ukazovateľov a alokácie haldy na jeden uzol. Konštrukcia umiestňuje každý uzol tak, aby jeho deti padli na voľné sloty v overovacom poli, čo umožňuje viacerým rodičom zdieľať oblasti, pokiaľ sa ich deti líšia v nasledujúcom bajte.

Pre slovník XLM-RoBERTa s 250 tisíc položkami sa celá štruktúra trie zmestí do ~9 MB súvislej pamäte (~780k slotov × 12 bajtov pre posun, overenie a ID tokenu na slot). Horúca pracovná súprava jedného prechodu sa pohybuje v ráde 100 KB, čo sa pohodlne zmestí do L2, pričom hardvérový prednačítač sa prispôsobí pravidelnému vzoru prístupu.

Referencia naopak udržiava hashovaciu mapu detí pre každý uzol, čo stojí približne štyri závislé načítania na jeden bajtový krok. Hashuje sa bajt, načíta sa vedro, nasleduje sa ukazovateľ na položku vedra, nasleduje sa ďalší ukazovateľ na detský uzol v halde a nakoniec sa konzultuje vlastná hashovacia mapa pre ďalší bajt. Objekty v halde sú rozptýlené a väčšina načítaní minie vyrovnávaciu pamäť.

SP a IREE ukladajú ID tokenu priamo v uzle trie. Referenčná implementácia trie vracia iba zhodné bajty, takže encoder musí alokovať nový reťazec pre každú zhodu, vyhľadať ho v bočnej hashovacej mape na získanie ID tokenu a prečítať skóre zo samostatného vektora. SP vracia ID ako uloženú hodnotu trie a číta skóre z plochého poľa s desatinnou čiarkou:

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žíva stavové rozhranie kurzora namiesto vracania všetkých zhôd v jedinom volaní, ale model nákladov zostáva rovnaký: ID tokenov a skóre sú dostupné ako indexy plochého poľa, bez alokácie alebo hashovania pri samotnom prechode trie. IREE alokuje iba jeden výstupný buffer na jeden prechod; SP navyše alokuje uzol mriežky na zhodu počas dynamického programovania, čo predstavuje celkovo približne 1 559 alokácií pri 514 tokenoch.

Pred zmenou trie izolujeme náklady na algoritmus od nákladov na implementáciu: ako rýchlo beží referencia po odstránení zbytočnej práce?

Základná línia zachováva referenčnú štruktúru trie s deťmi v hashovacej mape, ale vykonáva dve zmeny. Dočasná štruktúra vlastnená volajúcim uchováva tabuľku DP, predbežne spracované bajty a ID výstupných tokenov, takže kritická cesta nevykonáva po zahriatí žiadne alokácie. ID tokenu a skóre sú uložené priamo v uzle trie, čím sa odstraňuje konverzia reťazcov na zhodu a vyhľadávanie v bočnej hashovacej mape.

Vnútorná sluč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,
        }
    }
}

Ide o rovnaký algoritmus s rovnakým tvarom trie a rovnakým prístupovým vzorom, len bez alokácií a redundantných vyhľadávaní. Výsledky pri 514 tokenoch:

Engine	p50	Inštrukcie	IPC	Alokácie
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
Základná línia	156 µs	1,53M	2,62	0

Základný model bez alokácií v porovnaní s existujúcimi implementáciami tokenizátora pri 514 tokenoch.

Odstránenie materializácie reťazcov na zhodu a sekundárneho vyhľadávania v hashovacej mape znižuje počet inštrukcií 2,4-krát v porovnaní s referenciou. Hodnota IPC s nimi klesá (2,86 → 2,62), pretože hľadanie ukazovateľa v hashovacej mape, predtým skryté za prácou alokátora, je teraz dominantnou cenou na jeden bajt.

Referenčný tokenizátor a základný tokenizátor bez alokácií p50 a p99 latencia podľa veľkosti vstupu.

Miera zlyhaní L2 cache referenčného tokenizátora a základného tokenizátora bez alokácií podľa veľkosti vstupu.

Základná línia znižuje latenciu p50 2- až 3-násobne v porovnaní s referenciou pri každej veľkosti a miera minutia L2 zostáva blízko 2–3 % do 2k tokenov oproti 8–48 % u referencie. Zostávajúce náklady predstavuje hľadanie ukazovateľa v štruktúre HashMap, ktorým sa zaoberá nasledujúca časť.

Trie založená na hashovacej mape predstavuje hľadanie ukazovateľa. SP a IREE to riešia pomocou dvojpoľovej štruktúry trie (pozri predchádzajúcu časť) a my si túto dátovú štruktúru požičiavame. Rozdiel spočíva v tom, že SP a IREE sú všeobecné knižnice podporujúce streamovaný vstup, výstupné posuny a viaceré režimy kódovania; táto flexibilita ich stojí dodatočnú správu okolo trie (mriežku v SP, osemstupňový kanál v IREE). My nič z toho nepotrebujeme, takže trie vkladáme priamo do Viterbiho slučky a zvyšok vynechávame.

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

Bajtový krok stojí dve čítania z poľa, celočíselné sčítanie a porovnanie. Celá štruktúra trie má približne 9 MB, ale každý jednotlivý prechod prechádza iba cez ~100 KB z nej, takže horúce dáta zostávajú v L2.

Engine	p50	p99	Inštrukcie	Minutie L2 %	Minutie 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 %
Základná línia	155 µs	1 170 µs	1,53M	2,59 %	16,66 %
+ Darts	68 µs	80 µs	1,10M	6,45 %	11,42 %

Latencia tokenizátora, inštrukcie a miera zlyhaní cache po zavedení double-array trie.

Reálny čas klesá 2,3-násobne oproti základnej línii a 4,8-násobne oproti referencii. Inštrukcie klesajú o ďalších 28 % oproti základnej línii. Odstránenie mriežky a kanála nám umožňuje vykonať 1,10 milióna inštrukcií na jeden prechod oproti 1,82 milióna v SP a 1,61 milióna v IREE pri rovnakom tvare trie; rozdiel predstavuje administratívna réžia, ktorú udržiavajú pre všeobecné rozhrania API, ktoré nepotrebujeme.

Štruktúra Darts trie z optimalizácie 1 vyžaduje dve závislé načítania na jeden bajtový krok. Najprv základný posun rodiča, potom záznam check v slote kandidátskeho dieťaťa na potvrdenie, že ho rodič vlastní. Záznam check slúži iba na overenie vlastníctva. Na tú istú otázku možno odpovedať pomocou bitmapy pre každý uzol, ktorá zaznamenáva, ktoré bajty majú deti. Z konštrukcie dvojpoľovej trie vyplýva, že bit pre bajt b v mape rodiča je nastavený presne vtedy, keď slot kandidáta pre daný bajt vlastní rodič. Bitmapa a pole check nesú rovnakú informáciu, takže pole check môžeme zo štruktúry za behu úplne odstrániť.

Vyhľadanie v bitmape sa preklada do jediného testu bitov voči jednému 64-bitovému slovu, čo prebehne približne za jeden cyklus. Súvisiaca inštrukcia POPCNT, ktorá počítajú všetky nastavené bity v registri, sa používa v iných návrhoch trie indexovaných bitmapou na priamy výpočet posunu dieťaťa z bitmapy. My ju nepotrebujeme, pretože konštrukcia Darts už tento posun ukladá v poli base uzla.

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

Keď už sme pri tom, zabalíme štyri polia na jeden uzol (bitmapu, base, ID tokenu a skóre) do jedného 64-bajtového riadku vyrovnávacej pamäte. Jediná návšteva uzla teraz stiahne všetko: skóre a ID tokenu, ktoré potrebujeme v koncových uzloch, plus bitmapu a base pre ďalší krok. Pole check sa používa iba počas konštrukcie a je odstránené zo štruktúry za behu. Indexovanie kritickej cesty je nekontrolované, pretože konštrukcia trie a štruktúra slučky už zaručujú dodržanie hraníc. Kompromisom je priestor: pri 64 bajtoch na uzol a ~780k slotoch sa štruktúra trie zväčšuje z ~9 MB na ~50 MB.

Výsledky pri 514 tokenoch:

Engine	p50	p99	Inštrukcie	Prístupy L2	Minutie L2 %	Minutie 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 %
Základná línia	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 inline balenie	65 µs	77 µs	1,06M	1,8K	20,53 %	15,01 %

Výsledky tokenizátora pri 514 tokenoch po pridaní bitmapy a inline balenia.

Reálny čas systematicky klesá o ďalších 4,5 % pri 514 tokenoch a rastie s dĺžkou vstupu, pričom dosahuje 6–7 % pri 4k a 16k. Odstránenie poľa check mení krok trie na jeden bajt z dvoch načítaní riadku vyrovnávacej pamäte na jedno, čo predstavuje 2,5-krát menej prístupov do L2 na jeden prechod (1,8K oproti 4,6K pri obyčajnej štruktúre Darts). Miera minutia vyskočí zo 6 % na 20 %, ale väčšinu z toho tvorí menovateľ: s menej ako polovičným počtom prístupov zaberajú minutia, ktoré máme, väčší podiel.

Dáta štruktúry trie zaberajú približne 50 MB (780 000 uzlov po 64 bajtov). V predvolenom procese systému Linux je táto pamäť zabezpečená 4 KB stránkami, takže celá štruktúra trie pokrýva približne 12 000 rôznych virtuálnych stránok, čo ďaleko presahuje kapacitu vyrovnávacej pamäte pre preklad adries prvej úrovne procesora (Sapphire Rapids ich má 96). Každý Viterbiho krok sa dotýka iného uzla a počas kódovania 512 tokenov spotrebujeme približne 9 000 cyklov na prechody tabuľkami stránok, čo predstavuje asi 3 % z rozpočtu na jeden prechod.

Riešením je podoprieť štruktúru trie 2 MB obrovskými stránkami namiesto 4 KB stránok. Tých istých 50 MB teraz pokrýva 25 stránok, zmestí sa kompletne do vyrovnávacej pamäte Translation Lookaside Buffer a samotné adresáre tabuliek stránok prestávajú konkurovať štruktúre trie o vyrovnávaciu pamäť L2.

Linux ponúka dva spôsoby, ako zabezpečiť pamäť pomocou 2 MB stránok. Oportunistický prístup spočíva v transparentných obrovských stránkach (THP): váš program požiada jadro o povýšenie oblasti pamäte zo 4 KB stránok na 2 MB stránky a vlákno jadra na pozadí sa to snaží vykonať, keď má čas a k dispozícii súvislú fyzickú pamäť. Druhou alternatívou je alokovať štruktúru trie priamo z vopred vyhradenej fondu obrovských stránok, ktorý jadro vyčlení pri zavádzaní systému:

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

Tento fond musí existovať na hostiteľskom systéme (vm.nr_hugepages nastavené pri štarte). V našich produkčných inštanciách je vyhradených 10 561 obrovských stránok; štruktúra trie z nich spotrebováva 24. Po funkcii mmap súbor /proc/<pid>/smaps potvrdzuje danú oblasť: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Zisk predstavuje 3–12 % reálneho času naprieč dĺžkami vstupov, čo je konzistentné a úmerné dĺžke pracovnej súpravy:

Tokeny	Bez obrovských stránok	S obrovský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 %

Latencia tokenizátora s veľkými stránkami aj bez nich naprieč veľkosťami vstupu.

Najväčší prínos sa dosahuje pri 4k tokenoch, kde je pracovná súprava dostatočne veľká na to, aby prevádzka tabuliek stránok aktívne konkurovala dátam trie o šírku pásma L2. Za hranicou 4k sa prínos zmenšuje, pretože minutia L3 dominujú a prechody tabuľkami stránok tvoria menší podiel z celkového času.

Po spojení týchto troch optimalizácií prinášame meranie finálneho enginu v porovnaní so všetkými ostatnými tokenizátormi RoBERTa, ktoré sme testovali. Okrem natívnych binárnych súborov C++ SentencePiece a C IREE sme otestovali aj hotové rustové obalové balíky, po ktorých by produkčný používateľ siahol, konkrétne sentencepiece a iree-tokenizer. Balík sentencepiece vykazuje 1,6- až 1,9-krát vyššiu latenciu než natívny binárny súbor C++ pri každej dĺžke vstupu, pričom réžia pochádza z toho, že obal alokuje nový zoznam častí tokenov, pričom každá obsahuje vlastnú kópiu textu časti pri každom prechode. Balík iree-tokenizer prispôsobuje veľkosť svojho interného buffera každému vstupu, čo eliminuje konštantnú alokáciu na jeden prechod, ktorú natívny binárny súbor IREE platí bez ohľadu na veľkosť vstupu, čo je dôležité najmä pri krátkych vstupoch, kde by inak alokácia dominovala skutočnej práci tokenizácie. Pri dĺžkach vstupov v produkcii sa konštantné náklady amortizujú a vlastná réžia obalu z kopírovania tokenov späť do jazyka Rust pri jednom volaní tlačí jeho latenciu nad úroveň natívneho binárneho súboru.

p50 a p99 latencia na jedno kódovanie naprieč implementáciami tokenizátora a veľkosťami vstupu.

Finálny engine spúšťa rovnaký Viterbiho algoritmus nad rovnakým slovníkom s 250 000 tokenmi, na kritickej ceste nič nealokuje a produkuje tokenovo presný výstup v porovnaní s referenciou. Každý bajtový krok teraz predstavuje jediné načítanie 64-bajtového riadku vyrovnávacej pamäte, jeden test bitmapy a jedno celočíselné sčítanie, čo predstavuje pokles oproti kontrole HashMap, volaniu alokátora a niekoľkým špeciálnym vetveniam v základnej implementácii. Inštrukcie na jeden prechod klesli z 3,66 milióna na 1,04 milióna, čo predstavuje 3,5-násobný pokles. Reálny čas zodpovedá tomuto pomeru pri krátkych vstupoch a ďalej sa zväčšuje pri dlhých vstupoch, kde alokácie referencie na jeden token začínajú pretekať do L2 a L3.

V praxi táto práca znížila využitie CPU v vrstve predbežného spracovania nášho inferenčného zásobníka 5- až 6-násobne a skrátila latenciu pre-rankera o dvojciferné milisekundy.

Väčšina zníženia latencie vyplynula z ľahkých zmien, ktoré odstránili zbytočnú prácu: vynechanie alokácií na jeden prechod, nahradenie štruktúry trie s deťmi v hashovacej mape dvojpoľovou štruktúrou a odstránenie réžie všeobecného rozhrania API. Využitie existujúcej bitmapy na jeden uzol ako kontroly platnosti a podopretie trie 2 MB obrovskými stránkami prispeli celkovo k 10–20 % navyše k väčšiemu prechodu na dvojpoľovú štruktúru trie, v závislosti od dĺžky vstupu.

Hotové rustové obalové balíky pre SentencePiece a IREE ilustrujú rovnaký bod opačne. Okrem natívnych tokenizátorov v C a C++ sú rozhranie FFI a alokácie na jedno volanie, ktoré obaly zavádzajú, viditeľné v číslach latencie. Zdrojový kód v jazyku Rust pre finálny Viterbiho encoder je k dispozícii 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}
}

Motivácia

Viterbiho algoritmus

Referenčná implementácia

Poznámka k predbežnému spracovaniu

Kam referenčná implementácia míňa svoj čas

Ako sme merali

Prečo sú SentencePiece a IREE rýchlejšie: Dvojpoľová štruktúra Trie

Základná línia: Nulovo-alokačný port referencie

Optimalizácia 1: Dvojpoľová štruktúra Trie bez dodatočných abstrakcií

Optimalizácia 2: Bitmapové a inline balenie

Optimalizácia 3: Obrovské stránky pre trie

Záverečné porovnanie

Záverečné myšlienky

Referencie

Citácia