Poboljšanje performansi CPU-a Unigram tokenizatora

Ponovno smo implementirali naš Unigram tokenizator od temelja kao ciljani projekt performansi. Pri ulaznim duljinama u produkciji, novi koder smanjuje p50 latenciju za otprilike 5× u odnosu na Hugging Face tokenizers crate, ~2× u odnosu na SentencePiece (C++) i ~1,5× u odnosu na IREE tokenizator (C), uz nula alokacija u gomili u stacionarnom stanju. U produkciji je ovo smanjilo iskorištenje CPU-a u našem stogu za inferenciju za 5-6× i skratilo latenciju rerankera za dvoznamenkasti broj milisekundi.

Trošak inferencije kod LLM-ova obično se prikazuje kao rad GPU-a: KV predmemorije, jezgre pažnje, usmjeravanje stručnjaka. No značajan dio produkcijskog prometa odvija se na znatno manjim modelima, dva do tri reda veličine manjima od graničnih transformera. Modeli za ugrađivanje (embedding), klasifikatori i rerankeri i dalje se izvode na GPU-ovima, ali prolaz naprijed (forward pass) dovoljno je kratak da sve oko njega počinje biti važno.

Najjasniji primjer je reranker koji boduje stotine kandidata dokumenata po zahtjevu. S malim modelom, izračun na GPU-u obično završava u jednoznamenkastom broju milisekundi, ali svaki ulaz i dalje prvo mora proći kroz tokenizaciju na strani CPU-a, pretvarajući tekst u ID-jeve rječnika. Kada je serija velika, to može činiti značajan udio ukupne latencije zahtjeva.

Ovaj rad cilja XLM-RoBERTa, Unigram rječnik s 250 tisuća tokena obučen pomoću SentencePiecea. Naknadno ugođeni koderi obitelji RoBERTa uobičajen su izbor za zadatke rangiranja, pronalaženja i sličnosti. Ovdje navodimo latenciju po kodiranju s jednom dretvom; u produkciji radimo paralelizaciju preko jezgri pomoću rayona.

Tijekom ovog posta tokenizers crate od Hugging Facea tretiramo kao referencu s kojom mjerimo performanse, budući da je to zadani Rust tokenizator kojem se većina timova okreće. SentencePiece (C++) i IREE (C) uključeni su radi usporedbe kao druge dvije velike produkcijske implementacije.

Unigram tokenizacija, koju je uvrstio Kudo 2018. (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) i isporučio u SentencePiece, prikazuje segmentaciju kao problem najvjerojatnije putanje: svaki token rječnika ima naučenu log-vjerojatnost, a koder bira segmentaciju čiji zbroj ocjena tokena daje najveću vrijednost. Problem se rješava Viterbijevim algoritmom (1967.), s pozicijama bajtova kao slojevima grafa i tokenima rječnika kao bridovima koji obuhvaćaju uzastopni raspon bajtova. Rekurencija je jednodimenzionalno dinamičko programiranje preko pozicija bajtova:

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

Vanjska petlja je linearna u odnosu na duljinu ulaza, ograničena na svakoj poziciji najduljim tokenom rječnika (48 UTF-8 bajtova za XLM-RoBERTa). Unutrašnji prolaz kroz stablo prefiksa (trie) predstavlja kritičnu stazu (hot path): na ulazu od 16k tokena izvršava stotine tisuća prijelaza i završava desetke milijuna instrukcija po kodiranju.

Tokenizers v0.22 od Hugging Facea implementira rekurenciju izravno, s tablicom raspršenog adresiranja (hash-map) za obrnuto traženje i djecom tablice raspršenog adresiranja u svakom čvoru stabla:

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
}

I unutarnja petlja (pojednostavljena) – svako podudaranje pretrage prefiksa alocira novi niz, a zatim ispituje tablicu raspršenog adresiranja:

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
}

Kod je čitljiv Rust, ali događa se puno posla koji je nevidljiv na razini izvornog koda. Sljedeći odjeljak razlaže kamo odlazi vrijeme.

Prije nego što Viterbi započne s radom, ulaz prolazi kroz Unicode normalizaciju (NFKC) i označavanje granica riječi. Profiliranje pokazuje da predobrada čini manje od 5% ukupnog vremena kodiranja u našim scenarijima, stoga se u ostatku posta fokusiramo na Viterbijev prolaz naprijed.

Postoje dva načina da se ubrza CPU kod: izvršiti manje instrukcija ili postići viši IPC. Cachegrind mjeri prvo, perf stat mjeri drugo. Svaka optimizacija u ovom postu spada u jednu ili drugu kategoriju.

CPU-ovi zastajkuju zbog četiriju stvari: promašaja predmemorije (L1d ~4 ciklusa, promašaj L3 u DRAM broji se u stotinama), lanaca zavisnih učitavanja koje mehanizam za izvanredno izvršavanje (out-of-order) ne može paralelizirati, pogrešnih predviđanja grana (~15-20 ciklusa pražnjenja po pogrešnoj pretpostavci) i poziva alokatora. Referentna implementacija unigram mehanizma pogađa sve četiri:

Usko grlo	Mehanizam	Izmjereni utjecaj
Alokacija po podudaranju	String::from_utf8 nad podudarnim bajtovima plus pretraga u AHashMap za svako podudaranje u stablu; ID tokena se već nalazi u listu, ali referenca ga ondje ne pohranjuje	7.295 alokacija kod 512 tokena, 14.542 kod 1k, 299.171 kod 16k
Potraga za pokazivačima po bajtu	AHashMap na svakom čvoru stabla, uz raspršivanje (hashing) sljedećeg bajta i dereferenciranje dvaju pokazivača po koraku; struktura Node također grupira hladna polja poput is_leaf, uvlačeći ih u predmemoriju pri svakom posjetu	Stopa promašaja L1d ostaje niska (~2,5% kod 1k) jer se vruće stranice alokatora nalaze u L1; trošak se prikazuje kao latencija ovisnog učitavanja
Zagušenje L2 na dugačkim ulazima	DP tablica, privremeni prostor prefiksa i izlazna lista tokena nanovo se alociraju pri svakom pozivu	Stopa promašaja L2 penje se s 12% kod 128 tokena na 50% kod 16k kako kumulativni otisak alokacije premašuje L2 od 2 MB po jezgri

Referentna uska grla tokenizatora, mehanizmi i izmjereni učinci.

Glavni trošak je alokacija po podudaranju. Pri 1k tokena, referentna implementacija izvršava (dovršava) otprilike 7 milijuna instrukcija po kodiranju uz IPC od 2,8, od čega većinu u pozivima alokatora i pretvorbama nizova. Promašaji L2 važni su tek nakon ~1k tokena, kada kumulativni otisak alokacija preplavi predmemoriju po jezgri.

Latencija referentnog tokenizatora p50 i p99 prema veličini ulaza.

Stopa promašaja L2 predmemorije referentnog tokenizatora prema veličini ulaza.

Bajtovi alocirani po encodeu od strane referentnog tokenizatora kroz različite veličine ulaza.

Alokacija po tokenu je konstantna: oko 2 KB i ~18 alokacija po tokenu, bez obzira na veličinu ulaza. Nelinearnost latencije dolazi od L2: kumulativni bajtovi premašuju 2 MB oko 1k tokena, stopa promašaja penje se s 8% na ~50%, a krivulja p99 se savija u skladu s tim.

Uspoređujemo tri produkcijska tokenizatora, od kojih svi pokreću isti Unigram algoritam nad istim XLM-RoBERTa modelom:

Mehanizam	Jezik	Stablo
Hugging Face	Rust	Stablo prefiksa s djecom u HashMapu
SentencePiece	C++	Darts stablo s dvostrukim poljem
IREE tokenizator	C	Darts stablo s dvostrukim poljem

Implementacije tokenizatora, jezici i prikazi trie stabla korišteni u usporedbi.

SentencePiece i IREE povezujemo putem izravnog C/C++ FFI-ja umjesto njihovih omotača u Rustu; omotači dodaju alokacije po pozivu i pretvorbe nizova koje nisu dio troška temeljne biblioteke.

Svi mehanizmi izgrađeni su uz potpunu optimizaciju (O3, ciljanje nativnog sklopa, optimizacija vremena povezivanja). Mjerenja se izvode vezano uz jednu tihu jezgri na procesoru Intel Xeon Platinum 8488C. Za svaki mehanizam i veličinu ulaza prikupljamo latenciju zidnog sata (p50 i p99, 10.000 iteracija nakon 1.000 zagrijavanja), izvršene instrukcije putem Valgrind cachegrind alata, IPC i stope promašaja predmemorije iz hardverskih brojača perf stat, te alokacije u gomili koje broji omotač globalnog alokatora oko mjerenog prozora.

Pri 514 tokena (+2 preko 512 predstavlja BOS/EOS injekciju):

Implementacija	p50 latencija	Instrukcije	IPC	Alokacije
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

Latencija tokenizatora, broj instrukcija, IPC i alokacije pri 514 tokena.

SP i IREE postižu otprilike trećinu latencije referentne implementacije. Izvršavaju upola manje instrukcija, rade uz viši IPC i alociraju jedan do dva reda veličine manje.

SentencePiece i IREE koriste stablo s dvostrukim poljem koje je uveo Aoe 1989. (Aoe, 1989; Karoonboonyanan). Stablo je kodirano u dva ravna cjelobrojna polja, base i check:

cpp

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

Pretraga djeteta je:

cpp

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

Svaki korak prolaska kroz stablo zahtijeva samo dva učitavanja iz polja, jedno cjelobrojno zbrajanje i jednu usporedbu, bez raspršivanja, bez potrage za pokazivačima i bez alokacija u gomili po čvoru. Konstrukcija smješta svaki čvor tako da njegova djeca padaju na slobodna mjesta u polju za provjeru valjanosti, dopuštajući višestrukim roditeljima da dijele regije sve dok se njihova djeca razlikuju po bajtu koji slijede.

Za rječnik od 250k tokena modela XLM-RoBERTa, cijelo stablo stane u ~9 MB susメモernog memorijskog prostora (~780k mjesta × 12 bajtova za pomak, provjeru valjanosti i ID tokena po mjestu). Vrući radni skup jednog kodiranja je reda veličine 100 KB, udobno unutar L2, a hardverski prefetcher zaključava se na redoviti uzorak pristupa.

Nasuprot tomu, referentna implementacija drži tablicu raspršenog adresiranja djece po čvoru, što košta otprilike četiri zavisna učitavanja po koraku bajta. Rasprši se bajt, učita se spremnik (bucket), prati se pokazivač do unosa u spremniku, prati se drugi pokazivač do čvora djeteta na gomili, a zatim se provjerava njegova vlastita tablica raspršenog adresiranja za sljedeći bajt. Objekti na gomili su raspršeni i većina učitavanja promašuje predmemoriju.

SP i IREE pohranjuju ID tokena izravno u čvor stabla. Referentna implementacija stabla vraća samo podudarne bajtove, stoga koder mora alocirati novi niz za svako podudaranje, potražiti ga u pomoćnoj tablici raspršenog adresiranja kako bi vratio ID tokena i pročitati ocjenu iz zasebnog vektora. SP vraća ID kao pohranjenu vrijednost stabla i čita ocjenu iz ravnog polja s pomičnim zarezom:

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 koristi API s upravljanjem stanjem (stateful cursor API) umjesto vraćanja svih podudaranja u jednom pozivu, ali model troškova je isti: ID-jevi tokeni i ocjene dostupni su kao indeksi u ravnom polju, bez alokacije ili raspršivanja pri samom prolasku kroz stablo. IREE alocira samo jedan izlazni spremnik po kodiranju; SP dodatno alocira čvor rešetke po podudaranju tijekom DP-a, što ukupno iznosi otprilike 1.559 alokacija pri 514 tokena.

Prije mijenjanja stabla, odvajamo trošak algoritma od troška implementacije: koliko brzo referentna implementacija radi s uklonjenim suvišnim poslom?

Osnovna linija zadržava stablo s djecom u tablici raspršenog adresiranja iz referentne implementacije, ali uvodi dvije promjene. Pomoćna struktura u vlasništvu pozivatelja drži DP tablicu, predobrađene bajtove i izlazne ID-jeve tokena, tako da kritična staza ne obavlja nikakve alokacije nakon zagrijavanja. ID tokena i ocjena pohranjuju se izravno u čvor stabla, čime se uklanjaju pretvorba nizova po podudaranju i pretraga u pomoćnoj tablici raspršenog adresiranja.

Unutarnja petlja:

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

Ovo je isti algoritam, s istim oblikom stabla i istim uzorkom pristupa, samo bez alokacija i bez suvišnih pretraga. Rezultati pri 514 tokena:

Mehanizam	p50	Instrukcije	IPC	Alokacije
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
Osnovna linija	156 µs	1,53M	2,62	0

Početna vrijednost bez alokacija u usporedbi s postojećim implementacijama tokenizatora pri 514 tokena.

Uklanjanje materijalizacije nizova po podudaranju i sekundarne pretrage u tablici raspršenog adresiranja smanjuje broj instrukcija 2,4× u odnosu na referentnu implementaciju. IPC pada s njima (2,86 → 2,62) jer potraga za pokazivačima u tablici raspršenog adresiranja, prethodno skrivena iza rada alokatora, sada postaje dominantan trošak po bajtu.

Latencija referentnog i početnog tokenizatora bez alokacija p50 i p99 prema veličini ulaza.

Stopa promašaja L2 predmemorije referentnog i početnog tokenizatora bez alokacija prema veličini ulaza.

Osnovna linija smanjuje p50 latenciju za 2–3× u odnosu na referentnu implementaciju pri svakoj veličini, a stopa promašaja L2 ostaje blizu 2–3% do 2k tokena, u usporedbi s 8–48% kod referentne implementacije. Preostali trošak je potraga za pokazivačima u HashMapu, koju rješava sljedeći odjeljak.

Stablo s tablicom raspršenog adresiranja predstavlja potragu za pokazivačima. SP i IREE rješavaju ovaj problem stablom s dvostrukim poljem (vidi prethodni odjeljak), i mi posuđujemo istu strukturu podataka. Razlika je u tome što su SP i IREE općenite biblioteke koje podržavaju strujanje (streaming) ulaza, izlazne pomake i višestruke načine kodiranja; ta fleksibilnost ih košta dodatnog vođenja evidencije oko stabla (rešetka u SP-u, osamstupanjski cjevovod u IREE-u). Nama ništa od toga nije potrebno, pa stablo ugrađujemo izravno u Viterbijevu petlju i odbacujemo ostatak.

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

Korak po bajtu košta dva čitanja iz polja, cjelobrojno zbrajanje i usporedbu. Cijelo stablo zauzima oko 9 MB, ali bilo koje pojedinačno kodiranje prolazi kroz samo ~100 KB tog prostora, tako da vrući podaci ostaju u L2.

Mehanizam	p50	p99	Instrukcije	L2 promašaj %	L3 promašaj %
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%
Osnovna linija	155 µs	1.170 µs	1,53M	2,59%	16,66%
+ Darts	68 µs	80 µs	1,10M	6,45%	11,42%

Latencija tokenizatora, instrukcije i stope promašaja predmemorije nakon usvajanja dvostruko-niznog trie stabla.

Latencija zidnog sata pada 2,3× u odnosu na osnovnu liniju i 4,8× u odnosu na referentnu implementaciju. Instrukcije se smanjuju za dodatnih 28% u odnosu na osnovnu liniju. Uklanjanje rešetke i cjevovoda omogućuje nam izvršavanje 1,10 milijuna instrukcija po kodiranju u usporedbi s 1,82M kod SP-a i 1,61M kod IREE-ja pri istom obliku stabla; razlika je u evidenciji koju oni vode za općenite API-je koji nama nisu potrebni.

Darts stablo iz Optimizacije 1 zahtijeva dva zavisna učitavanja po koraku bajta. Prvo bazni pomak roditelja, zatim unos provjere (check) na mjestu kandidata za dijete kako bi se potvrdilo da je roditelj njegov vlasnik. Unos provjere služi samo za potvrdu vlasništva. Isto se pitanje može riješiti slikovnom mapom (bitmapom) po čvoru koja bilježi koji bajtovi imaju djecu. Prema konstrukciji stabla s dvostrukim poljem, bit za bajt b u bitmapi roditelja postavljen je točno onda kada je mjesto kandidata za taj bajt u vlasništvu roditelja. Bitmapa i polje provjere nose istu informaciju, tako da polje provjere možemo u potpunosti ukloniti iz izvođene strukture.

Pretraga u bitmapi kompajlira se u jednu provjeru bita nad jednom 64-bitnom riječi, što se izvršava u otprilike jednom ciklusu. Srodna instrukcija POPCNT, koja broji sve postavljene bitove u registru, koristi se u drugim dizajn E stabala indeksiranih bitmapom za izravan izračun pomaka djeteta iz bitmape. Nama to nije potrebno jer Darts konstrukcija već pohranjuje taj pomak u baznom polju čvora.

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

Kad smo već kod toga, pakiramo četiri polja po čvoru (bitmapu, bazu, ID tokena i ocjenu) u jedan predmemorijski redak od 64 bajta. Pojedinačni posjet čvoru sada dohvaća sve odjednom: ocjenu i ID tokena koji su nam potrebni na završnim čvorovima, plus bitmapu i bazu za sljedeći korak. Polje provjere koristi se samo tijekom konstrukcije i izbačeno je iz izvođenog rasporeda. Indeksiranje na kritičnoj stazi nije provjereno jer konstrukcija stabla i struktura petlje već garantiraju granice. Kompromis je prostor: s 64 bajta po čvoru × ~780k mjesta, stablo raste s ~9 MB na ~50 MB.

Rezultati pri 514 tokena:

Mehanizam	p50	p99	Instrukcije	L2 pristupi	L2 promašaj %	L3 promašaj %
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%
Osnovna linija	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 i ugrađeno pakiranje	65 µs	77 µs	1,06M	1,8K	20,53%	15,01%

Rezultati tokenizatora pri 514 tokena nakon dodavanja bitmapa i inline pakiranja.

Latencija zidnog sata sustavno pada za dodatnih 4,5% pri 514 tokena i raste s duljinom ulaza, dosežući 6–7% pri 4k i 16k. Izbacivanje polja provjere pretvara korak stabla po bajtu s dvaju učitavanja predmemorijskog retka na jedno, što rezultira 2,5× manje L2 pristupa po kodiranju (1,8k u usporedbi s 4,6k kod običnog Darts stabla). Stopa promašaja skče s 6% na 20%, ali većina toga otpada na nazivnik: s manje od upola manje pristupa, promašaji koje imamo zauzimaju veći udio.

Podaci stabla zauzimaju oko 50 MB (780.000 čvorova po 64 bajta). Na zadanom Linux procesu ta je memorija podržana stranicama od 4 KB, stoga cijelo stablo obuhvaća oko 12.000 različitih virtualnih stranica, što je daleko više nego što međuspremnik za translaciju pogleda (TLB) prve razine procesora može držati (Sapphire Rapids ima 96 unosa). Svaki Viterbi korak dodiruje drugi čvor, i tijekom kodiranja od 512 tokena potrošimo otprilike 9.000 ciklusa u obilascima tablice stranica, što čini oko 3% proračuna po kodiranju.

Rješenje je podržati stablo s ogromnim stranicama (huge pages) od 2 MB umjesto onih od 4 KB. Istih 50 MB sada obuhvaća 25 stranica, u potpunosti stane u TLB, a same tablične mape stranica prestaju se natjecati sa stablom za L2 predmemoriju.

Linux ima dva načina za podržavanje memorije stranicama od 2 MB. Oportunistički način su prozirne ogromne stranice (transparent huge pages – THP): vaš program traži od jezgre da nadogradi memorijsku regiju s 4 KB na stranice od 2 MB, a pozadinska nit jezgre pokušava to učiniti kada ima vremena i dostupne susjedne fizičke memorije. Druga mogućnost je izravna alokacija stabla iz unaprijed rezerviranog bazena ogromnih stranica koji jezgra izdvaja pri pokretanju sustava:

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

Bazen mora postojati na glavnom računalu (vm.nr_hugepages postavljeno pri pokretanju). Na našim produkcijskim instancama rezervirana je 10.561 ogromna stranica; stablo troši 24 njih. Nakon poziva mmap, /proc/<pid>/smaps potvrđuje regiju: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Dobitak iznosi 3–12% u latenciji zidnog sata kroz različite duljine ulaza, dosljedno je i proporcionalno duljini radnog skupa:

Tokeni	Bez ogromnih stranica	S ogromnim stranicama	Δ
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%

Latencija tokenizatora s huge pages i bez njih za različite veličine ulaza.

Najveći dobitak je pri 4k tokena, gdje je radni skup dovoljno velik da se promet tablice stranica aktivno natjecao s podacima stabla za propusnost L2. Iznad 4k dobitak se smanjuje jer L3 promašaji dominiraju, a obilasci tablice stranica postaju manji udio ukupnog vremena.

Spajajući tri optimizacije, ovdje je konačni mehanizam mjeren u usporedbi sa svakim drugim RoBERTa tokenizatorom koji smo testirali. Uz izvorne C++ SentencePiece i C IREE binarne datoteke, testirali smo i gotove Rust omotače (crates) koje bi produkcijski korisnik odabrao, sentencepiece i iree-tokenizer. Krate sentencepiece bilježi 1,6–1,9× višu latenciju od izvorne C++ binarne datoteke pri svakoj veličini ulaza, s opterećenjem koje dolazi od omotača koji alocira novu listu komadića tokena, od kojih svaki drži vlastitu kopiju teksta komadića, pri svakom kodiranju. Krate iree-tokenizer prilagođava veličinu svog internog spremnika svakom ulazu, što izbjegava stalnu alokaciju po kodiranju koju izvorna IREE binarna datoteka plaća bez obzira na veličinu ulaza, što je najvažnije kod kratkih ulaza gdje bi alokacija inače dominirala stvarnim radom tokenizacije. Pri produkcijskim duljinama ulaza stalni trošak se amortizira, a vlastito opterećenje omotača po pozivu od kopiranja tokena natrag u Rust gura njegovu latenciju iznad latencije izvorne binarne datoteke.

Latencija p50 i p99 po encodeu kroz implementacije tokenizatora i veličine ulaza.

Konačni mehanizam pokreće isti Viterbi nad istim rječnikom od 250.000 tokena, ne alocira ništa na kritičnoj stazi i proizvodi izlaz s točno jednakim tokenima u usporedbi s referentnom implementacijom. Svaki korak bajta sada je jedno učitavanje iz predmemorijskog retka od 64 bajta, jedna provjera bitmape i jedno cjelobrojno zbrajanje, što je smanjenje u odnosu na pretragu u HashMapu, poziv alokatora i nekoliko grana s posebnim slučajevima u osnovnoj implementaciji. Instrukcije po kodiranju pale su s 3,66M na 1,04M, što je pad od 3,5×. Latencija zidnog sata prati taj omjer kod kratkih ulaza i dodatno se povećava kod dugačkih ulaza, gdje alokacije referentne implementacije po tokenu počinju prelijevati L2 i L3.

U praksi je ovaj rad smanjio iskorištenje CPU-a u sloju predobrade našeg stoga za inferenciju za 5-6× i skratio latenciju rerankera za dvoznamenkasti broj milisekundi.

Glavni dio smanjenja latencije došao je od laganih promjena koje su uklonile suvišan posao: odbacivanja alokacija po kodiranju, zamjene stabla s djecom u HashMapu s dvostrukim poljem i uklanjanja opterećenja generičkog API-ja. Korištenje postojeće bitmape po čvoru kao provjere valjanosti i podržavanje stabla s ogromnim stranicama od 2 MB kombinirano su donijeli 10-20% poboljšanja povrh veće promjene na stablo s dvostrukim poljem, ovisno o veličini ulaza.

Gotove Rust omotačke krate oko SentencePiecea i IREE-ja ilustriraju istu poantu u obrnutom smjeru. Povrh izvornih C i C++ tokenizatora, FFI granica i alokacije po pozivu koje omotači uvode vidljive su u brojkama latencije. Izvorni kod u Rustu za konačni Viterbi koder dostupan je na 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}
}

Motivacija

Viterbijev algoritam

Referentna implementacija

Napomena o predobradi

Gdje referentna implementacija troši svoje vrijeme

Kako smo mjerili

Zašto su SentencePiece i IREE brži: Stablo s dvostrukim poljem

Osnovna linija (Baseline): Port referentne implementacije bez alokacija

Optimizacija 1: Stablo s dvostrukim poljem bez dodatnih apstraktnih slojeva

Optimizacija 2: Bitmapa i ugrađeno pakiranje

Optimizacija 3: Velike stranice za stablo

Konačna usporedba

Završne misli

Reference

Citiranje