Unapređenje performansi CPU-a Unigram tokenajzera

Ponovo smo implementirali naš Unigram tokenajzer od osnove kao usmereni projekat performansi. Pri ulaznim dužinama u produkciji, novi enkoder smanjuje p50 latenciju za otprilike 5× u poređenju sa Hugging Face tokenizers crate-om, ~2× u poređenju sa SentencePiece-om (C++) i ~1.5× u poređenju sa IREE-ovim tokenajzerom (C), uz nula heap alokacija u stacionarnom stanju. U produkciji je ovo smanjilo iskorišćenost CPU-a u našem inferencijskom steku za 5-6× i skratilo latenciju rerankera za dvocifreni broj milisekundi.

Troškovi LLM inferencije se obično posmatraju kroz rad GPU-a: KV keševi, attention jezgra, rutiranje eksperata. Međutim, značajan deo produkcijskog saobraćaja se izvršava na znatno manjim modelima, koji su dva do tri reda veličine manji od srodnih (frontier) transformera. Modeli za embedding, klasifikatori i rerankeri se i dalje izvršavaju na GPU-ovima, ali je prolazni prolaz (forward pass) dovoljno kratak da sve oko njega počinje da bude važno.

Najjasniji primer je reranker koji boduje stotine dokumenta kandidata po zahtevu. Sa malim modelom, GPU izračunavanje se obično završava u roku od jednocifrenog broja milisekundi, ali svaki ulaz i dalje mora najpre da prođe kroz tokenizaciju na strani CPU-a, čime se tekst pretvara u ID-ijeve rečnika. Kada je serija (batch) velika, to može da čini značajan udeo ukupne latencije zahteva.

Ovaj rad je usmeren na XLM-RoBERTa, Unigram rečnik veličine 250 hiljada tokena obučen pomoću alata SentencePiece. Fino podešeni enkoderi iz porodice RoBERTa predstavljaju čest izbor za zadatke rangiranja, pretrage i sličnosti. U celom tekstu navodimo latenciju po jednoj niti po enkodu; u produkciji vršimo paralelizaciju na više jezgara koristeći rayon.

Tokom ovog posta tretiramo Hugging Face-ov tokenizers crate kao referencu sa kojom vršimo upoređivanje performansi, s obzirom na to da je to podrazumevani Rust tokenajzer za koji se većina timova opredeljuje. SentencePiece (C++) i IREE (C) su uključeni radi poređenja kao druge dve vodeće produkcijske implementacije.

Unigram tokenizacija, koju je predstavio Kudo 2018. godine (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) i koja je implementirana u SentencePiece, posmatra segmentaciju kao problem putanje sa najvećom verovatnoćom: svaki token rečnika ima naučenu log-verovatnoću, a enkoder bira onu segmentaciju čiji zbir bodova tokena daje najveću vrednost. Problem se rešava pomoću Viterbijevog algoritma (1967), pri čemu pozicije bajtova predstavljaju slojeve grafa, a tokeni rečnika predstavljaju grane koje pokrivaju neprekidan opseg bajtova. Rekurzija predstavlja jednodimenzionalno dinamičko programiranje (DP) nad pozicijama 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

Spoljašnja petlja je linearna u odnosu na dužinu ulaza, ograničena na svakoj poziciji najdužim tokenom rečnika (48 UTF-8 bajtova za XLM-RoBERTa). Unutrašnji prolaz kroz stablo (trie walk) predstavlja kritičnu putanju (hot path): na ulazu veličine 16k tokena, on izvršava stotine hiljada prelaza i obavlja desetine miliona instrukcija po enkodu.

Hugging Face-ov tokenizers v0.22 direktno implementira rekurziju, uz korišćenje heš-mape (hash-map) za obrnuto pronalaženje i heš-mapu dece 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 unutrašnja petlja (pojednostavljena) – svako poklapanje pri pretrazi prefiksa alocira novu nisku (string), nakon čega vrši proveru heš-mape:

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 čitak Rust, ali se dešava mnogo radnji koje nisu vidljive na nivou izvornog koda. Sledeći odeljak detaljno razmatra na šta se vreme troši.

Pre nego što Viterbijev algoritam počne sa izvršavanjem, ulaz prolazi kroz Unicode normalizaciju (NFKC) i označavanje granica reči. Profilisanje pokazuje da pretprocesiranje čini manje od 5% ukupnog vremena enkodiranja u našim scenarijima, tako da se u ostatku teksta fokusiramo na Viterbijev proslérni prolaz (forward pass).

Postoje dva načina da se ubrza CPU kod: izvršavanje manjeg broja instrukcija ili postizanje veće vrednosti IPC (instrukcija po ciklusu). Cachegrind meri prvo, dok perf stat meri drugo. Svaka optimizacija u ovom postu se odnosi na jedno ili drugo.

Procesori zastaju zbog četiri stvari: promašaja keša (L1d ~4 ciklusa, promašaj L3 keša ka DRAM-u traje stotine ciklusa), lanaca zavisnih učitavanja (dependent-load chains) koje vanredosni (out-of-order) mehanizam ne može da paralelizuje, pogrešnih predviđanja grananja (~15–20 ciklusa se poništava po pogrešnoj proceni) i poziva alokatora. Referentna implementacija unigram enkodera nailazi na sve četiri prepreke:

Usko grlo	Mehanizam	Izmereni uticaj
Alokacija po poklapanju	String::from_utf8 nad poklopljenim bajtovima uz AHashMap pretragu po svakom poklapanju u stablu; ID tokena se već nalazi u listu, ali ga referentna implementacija tu ne čuva	7.295 alokacija na 512 tokena, 14.542 na 1k, 299.171 na 16k
Jurenje pokazivača po bajtu	AHashMap u svakom čvoru stabla, koja hešira sledeći bajt i dereferencira dva pokazivača po koraku; Node struktura takođe objedinjuje retko korišćena polja (cold fields) kao što je is_leaf, uvlačeći ih u keš pri svakoj poseti	Stopa promašaja L1d keša ostaje niska (~2,5% na 1k) zato što se aktivne stranice alokatora nalaze u L1 kešu; trošak se iskazuje kao latencija zavisnog učitavanja
L2 keš trashing (zagušenje) na dugim ulazima	DP tabela, privremeni prostor (scratch) za prefikse i izlazna lista tokena se nanovo alociraju pri svakom pozivu	Stopa promašaja L2 keša raste sa 12% na 128 tokena do 50% na 16k, pošto kumulativni otisak alokacija premašuje L2 keš od 2 MB po jezgru

Referentne uska grla tokenizera, mehanizmi i izmereni uticaji.

Dominantan trošak predstavlja alokacija po poklapanju. Na 1k tokena, referentna implementacija izvršava (završava) otprilike 7 miliona instrukcija po enkodu uz IPC od 2,8, pri čemu većina otpada na pozive alokatora i konverzije niski. L2 promašaji postaju bitni tek nakon ~1k tokena, kada kumulativni otisak alokacija preoptereti keš po jezgru.

Referentni tokenizer p50 i p99 latencija prema veličini ulaza.

Stopa promašaja L2 keša referentnog tokenizera prema veličini ulaza.

Broj bajtova alociranih po kodiranju od strane referentnog tokenizera 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 potiče od L2 keša: kumulativni bajtovi prelaze granicu od 2 MB oko 1k tokena, stopa promašaja raste sa 8% na ~50%, a kriva p99 latencije se savija u skladu sa tim.

Upoređujemo tri produkcijska tokenajzera, pri čemu svi izvršavaju isti Unigram algoritam nad istim XLM-RoBERTa modelom:

Mehanizam (Engine)	Jezik	Stablo (Trie)
Hugging Face	Rust	Prefiksno stablo sa heš-mapom dece (HashMap-children prefix trie)
SentencePiece	C++	Darts stablo sa dvostrukim nizom (Darts double-array trie)
IREE tokenajzer	C	Darts stablo sa dvostrukim nizom (Darts double-array trie)

Implementacije tokenizera, jezici i trie reprezentacije korišćene u poređenju.

Povezujemo SentencePiece i IREE kroz direktan C/C++ FFI umesto njihovih Rust omotača (wrappers); omotači dodaju alokacije po pozivu i konverzije niski koje nisu deo troškova osnovne biblioteke.

Svi mehanizmi su izgrađeni uz pune optimizacije (-O3, ciljanje nativnog hardvera, optimizacija tokom povezivanja - LTO). Merenja su vršena uz fiksiranje na jedno mirno jegro na procesoru Intel Xeon Platinum 8488C. Za svaki mehanizam i veličinu ulaza prikupljamo latenciju merenu zidnim časovnikom (p50 i p99, 10.000 iteracija nakon 1.000 zagrevanja), izvršene instrukcije putem alata Valgrind cachegrind, IPC i stope promašaja keša preko perf stat hardverskih brojača, kao i heap alokacije koje meri omot globalnog alokatora oko posmatranog prozora.

Na 514 tokena (+2 preko 512 se odnosi na ubacivanje BOS/EOS tokena):

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 tokenizera, broj instrukcija, IPC i alokacije pri 514 tokena.

SP i IREE ostvaruju približno jednu trećinu latencije referentne implementacije. Oni izvršavaju upola manje instrukcija, rade pri višem IPC-u i alociraju jedan do dva reda veličine manje.

I SentencePiece i IREE koriste stablo sa dvostrukim nizom koje je uveo Aoe 1989. godine (Aoe, 1989; Karoonboonyanan). Stablo je kodovano u dva ravna celobrojna niza, 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 dece se vrši na sledeći način:

cpp

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

Svaki korak obilaska stabla zahteva samo dva učitavanja iz niza, jedno celobrojno sabiranje i jedno poređenje, bez heširanja, jurenja pokazivača i alokacija na heap-u po čvoru. Postupak konstrukcije smešta svaki čvor tako da njegova deca padaju na slobodna mesta u nizu provere, omogućavajući višestrukim roditeljima da dele regione sve dok se njihova deca razlikuju po bajtu koji sledi.

Za XLM-RoBERTa rečnik veličine 250k, celokupno stablo staje u ~9 MB neprekinute memorije (~780k mesta × 12 bajtova za oset (offset), proveru i ID tokena po mestu). Aktivni radni skup (working set) za jedno enkodiranje je reda veličine 100 KB, što komotno staje u L2 keš, dok hardverski prefetcher prepoznaje pravilan obrazac pristupa.

Referentna implementacija, suprotno tome, održava heš-mapu dece po čvoru, što košta otprilike četiri zavisna učitavanja po koraku bajta. Hešira se bajt, učitava se segment (bucket), prati se pokazivač do unosa u segmentu, prati se drugi pokazivač do čvora deteta na heap-u, i na kraju se konsultuje sopstvena heš-mapa za sledeći bajt. Objekti na heap-u su rasuti i većina učitavanja promašuje keš.

I SP i IREE čuvaju ID tokena direktno u čvoru stabla. Referentna implementacija stabla vraća samo poklopljene bajtove, zbog čega enkoder mora da alocira novu nisku za svako poklapanje, da je potraži u sporednoj heš-mapi kako bi povratio ID tokena, i da pročita skor iz posebnog vektora. SP vraća ID kao uskladištenu vrednost stabla i čita skor iz ravnog niza sa pokretnim 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 sa stanjem (stateful cursor API) umesto vraćanja svih poklapanja u jednom pozivu, ali je model troškova isti: ID-ijevi tokena i skorovi su dostupni kao indeksi u ravnom nizu, bez alokacija i heširanja tokom samog prolaza kroz stablo. IREE alocira samo jedan izlazni bafer po enkodu; SP dodatno alocira čvor rešetke (lattice node) po poklapanju tokom DP-a, što ukupno iznosi oko 1.559 alokacija na 514 tokena.

Pre nego što promenimo stablo, izdvajamo trošak algoritma od troška implementacije: kojom brzinom referentna implementacija radi kada se ukloni suvišan posao?

Osnovna linija zadržava stablo sa heš-mapom dece iz referentne implementacije, ali vrši dve izmene. Struktura privremenog prostora (scratch struct) u vlasništvu pozivaoca sadrži DP tabelu, pretprocesirane bajtove i izlazne ID-ijeve tokena, tako da kritična putanja ne vrši nikakve alokacije nakon zagrevanja. ID tokena i skor se čuvaju direktno u čvoru stabla, čime se uklanjaju konverzija niske po poklapanju i pretraga u sporednoj heš-mapi.

Unutrašnja 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, sa istim oblikom stabla i istim obrascem pristupa, samo bez alokacija i suvišnih pretraga. Rezultati na 514 tokena:

Mehanizam (Engine)	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 (Baseline)	156 µs	1,53M	2,62	0

Referentna verzija bez alokacija u poređenju sa postojećim implementacijama tokenizera pri 514 tokena.

Uklanjanje materijalizacije niske po poklapanju i sekundarne pretrage u heš-mapi smanjuje broj instrukcija za 2,4× u poređenju sa referentnom implementacijom. IPC opada sa njima (sa 2,86 na 2,62) zato što jurenje pokazivača u heš-mapi, koje je ranije bilo skriveno iza rada alokatora, sada predstavlja dominantan trošak po bajtu.

Referentni i bazni tokenizer bez alokacija p50 i p99 latencija prema veličini ulaza.

Referentni i bazni tokenizer bez alokacija stopa promašaja L2 keša 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 keša ostaje blizu 2–3% do 2k tokena, u poređenju sa 8–48% kod referentne implementacije. Preostali trošak predstavlja jurenje pokazivača u heš-mapi (HashMap), čime se bavi sledeći odeljak.

Stablo zasnovano na heš-mapi predstavlja jurenje pokazivača. SP i IREE rešavaju ovaj problem pomoću stabla sa dvostrukim nizom (pogledajte prethodni odeljak), i mi preuzimamo istu strukturu podataka. Razlika je u tome što su SP i IREE opšte namenske biblioteke koje podržavaju striming ulaz, izlazne opsege i više režimenih kodiranja; ta fleksibilnost ih košta dodatnog vođenja evidencije oko stabla (rešetka u SP-u, osmostepeni cevovod u IREE-u). Nama ništa od toga nije potrebno, tako da ugrađujemo stablo direktno 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 bajta košta dva čitanja iz niza, celobrojno sabiranje i poređenje. Celokupno stablo ima oko 9 MB, ali bilo koje pojedinačno enkodiranje prolazi kroz samo ~100 KB tog stabla, tako da aktivni podaci ostaju u L2 kešu.

Mehanizam (Engine)	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 (Baseline)	155 µs	1.170 µs	1,53M	2,59%	16,66%
+ Darts	68 µs	80 µs	1,10M	6,45%	11,42%

Latencija tokenizera, instrukcije i stope promašaja keša nakon usvajanja double-array trie strukture.

Latencija merena zidnim časovnikom opada 2,3× u odnosu na osnovnu liniju i 4,8× u odnosu na referentnu implementaciju. Broj instrukcija se smanjuje za još 28% u odnosu na osnovnu liniju. Uklanjanje rešetke i cevovoda nam omogućava da izvršimo 1,10M instrukcija po enkodu u poređenju sa 1,82M kod SP-a i 1,61M kod IREE-a pri istom obliku stabla; razlika potiče od mehanizama za praćenje koje oni održavaju za API-je opšte namene koji nama nisu potrebni.

Darts stablo iz Optimizacije 1 zahteva dva zavisna učitavanja po koraku bajta. Prvo bazni oset roditelja, a zatim unos provere (check) na mestu kandidata za dete kako bi se potvrdilo da je roditelj njegov vlasnik. Unos provere služi samo za potvrdu vlasništva. Na isto pitanje se može odgovoriti pomoću bit mape po čvoru koja beleži koji bajtovi imaju decu. Na osnovu konstrukcije stabla sa dvostrukim nizom, bit za bajt b u bit mapi roditelja je podešen tačno onda kada je kandidat mesto za taj bajt u vlasništvu roditelja. Bit mapa i niz proverenose iste informacije, tako da možemo u potpunosti da odbacimo niz provere iz strukture tokom izvršavanja.

Pretraga u bit mapi se kompajlira u jednu bitovnu proveru nad jednom 64-bitnom reči, što se izvršava za otprilike jedan ciklus. Srodna POPCNT instrukcija, koja broji sve postavljene bitove u registru, jeste ono što drugi dizajneri stabala indeksiranih bit mapom koriste da direktno izračunaju oset deteta iz bit mape. Nama to nije potrebno zato što Darts konstrukcija već čuva taj oset u polju osnove (base) č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)
}

Dok smo već kod toga, pakujemo četiri polja po čvoru (bit mapu, oset, ID tokena i skor) u jednu keš liniju veličine 64 bajta. Jedna poseta čvoru sada dobija sve odjednom: skor i ID tokena koji su nam potrebni u terminalnim čvorovima, plus bit mapu i oset za sledeći korak. Niz provere se koristi samo tokom konstrukcije i izbačen je iz rasporeda u memoriji tokom izvršavanja. Indeksiranje na kritičnoj putanji se obavlja bez provera (unchecked) zato što konstrukcija stabla i struktura petlje već garantuju granice. Kompromis je prostor: sa 64 bajta po čvoru × ~780k mesta, stablo raste sa ~9 MB na ~50 MB.

Rezultati na 514 tokena:

Mehanizam (Engine)	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 (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%
+ Bit mapa i ugrađeno pakovanje (+ Bitmap and Inline Packing)	65 µs	77 µs	1.06M	1,8K	20,53%	15,01%

Rezultati tokenizera pri 514 tokena nakon dodavanja bitmapa i inline pakovanja.

Latencija merena zidnim časovnikom sistematski opada za još 4,5% na 514 tokena i raste sa dužinom ulaza, dostižući 6–7% na 4k i 16k. Izbacivanje niza provere pretvara korak stabla po bajtu sa dva učitavanja keš linije u jedno, što predstavlja 2,5× manje L2 pristupa po enkodu (1,8K u poređenju sa 4,6K kod običnog Darts stabla). Stopa promašaja skace sa 6% na 20%, ali je većina toga posledica imenioca: sa manje od upola manje pristupa, promašaji koje imamo zauzimaju veći udeo.

Podaci stabla zauzimaju otprilike 50 MB (780.000 čvorova od po 64 bajta). U podrazumevanom Linux procesu, ta memorija je podržana stranicama veličine 4 KB, tako da puno stablo obuhvata oko 12.000 različitih virtuelnih stranica, što daleko premašuje ono što bafer za prevod adresa prvog nivoa procesora (Translation Lookaside Buffer - TLB) može da drži (Sapphire Rapids ima 96 unosa). Svaki korak Viterbijevog algoritma dodiruje drugačiji čvor, i tokom enkodiranja od 512 tokena trošimo otprilike 9.000 ciklusa na obilaske tabele stranica, što je oko 3% budžeta po enkodu.

Rešenje je da se stablo podrži ogromnim stranicama (huge pages) veličine 2 MB umesto onih od 4 KB. Istih 50 MB sada obuhvata 25 stranica, u potpunosti stane u TLB bafer, a same direktorijum tabele stranica prestaju da se takmiče sa stablom za L2 keš.

Linux ima dva načina da podrži memoriju sa stranicama od 2 MB. Oportunistički način su transparentne ogromne stranice (transparent huge pages - THP): vaš program traži od kernela da nadogradi memorijski region sa stranica od 4 KB na stranice od 2 MB, a pozadinska nit (background thread) kernela pokušava to da uradi kada ima vremena i dostupne neprekinute fizičke memorije. Druga alternativa jeste alokacija stabla direktno iz unapred rezervisanog basena ogromnih stranica koji kernel izdvaja prilikom pokretanja sistema:

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

Basen mora da postoji na domaćinu (vm.nr_hugepages podešeno pri pokretanju). Na našim produkcijskim instancama rezervisano je 10.561 ogromnih stranica; stablo troši 24 od njih. Nakon mmap sistemskog poziva, /proc/<pid>/smaps potvrđuje region: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Dobitak iznosi 3–12% u latenciji merenoj zidnim časovnikom u zavisnosti od veličine ulaza, što je dosledno i srazmerno dužini radnog skupa:

Tokeni (Tokens)	Bez ogromnih stranica	Sa 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 tokenizera sa i bez huge pages kroz različite veličine ulaza.

Najveći dobitak je ostvaren na 4k tokena, gde je radni skup dovoljno velik da se saobraćaj tabele stranica aktivno takmičio sa podacima stabla za propusni opseg L2 keša. Iznad 4k dobitak se smanjuje zato što L3 promašaji dominiraju, a obilasci tabele stranica postaju manji udeo u ukupnom vremenu.

Kada spojimo ove tri optimizacije, ovde je prikazan konačni mehanizam izmeren u poređenju sa svakim drugim RoBERTa tokenajzerom koji smo testirali. Pored izvornih C++ SentencePiece i C IREE binarnih datoteka, testirali smo i gotove Rust omotače (wrapper crates) koje bi produkcijski korisnik izabrao, sentencepiece i iree-tokenizer. Krate sentencepiece pokazuje 1,6–1,9× veću latenciju od izvorne C++ binarne datoteke pri svakoj veličini ulaza, pri čemu režijski trošak potiče od toga što omotač alocira novu listu delova tokena, od kojih svaki drži svoju kopiju teksta dela, pri svakom enkodu. Krate iree-tokenizer prilagođava veličinu svog internog bafera svakom ulazu, čime izbegava konstantnu alokaciju po enkodu koju izvorna IREE binarna datoteka plaća bez obzira na veličinu ulaza, što je najvažnije kod kratkih ulaza gde bi alokacija inače dominirala nad stvarnim radom tokenizacije. Pri produkcijskim dužinama ulaza konstantan trošak se amortizuje, a sopstveni režijski trošak omotača po pozivu koji nastaje kopiranjem tokena nazad u Rust podiže njegovu latenciju iznad latencije izvorne binarne datoteke.

p50 i p99 latencija po kodiranju kroz implementacije tokenizera i veličine ulaza.

Konačni mehanizam izvršava isti Viterbijev algoritam nad istim rečnikom od 250.000 tokena, ne alocira ništa na kritičnoj putanji i daje izlaz koji je po tokenima identičan referentnom. Svaki korak bajta sada predstavlja jedno učitavanje keš linije od 64 bajta, jednu proveru bit mape i jedno celobrojno sabiranje, što je manje u odnosu na pretragu u heš-mapi, poziv alokatora i nekoliko grana za posebne slučajeve u implementaciji osnovne linije. Broj instrukcija po enkodu je opao sa 3,66M na 1,04M, što predstavlja smanjenje od 3,5×. Latencija merena zidnim časovnikom prati taj odnos kod kratkih ulaza, a dodatno se povećava kod dugih ulaza, gde alokacije po tokenu u referentnoj implementaciji počinju da preplavljuju L2 i L3 keš.

U praksi je ovaj rad smanjio iskorišćenost CPU-a u sloju pretprocesiranja našeg inferencijskog steka za 5-6× i skratio latenciju rerankera za dvocifreni broj milisekundi.

Glavni deo smanjenja latencije potiče od laganih izmena koje su uklonile suvišan rad: izbacivanje alokacija po enkodu, zamenjivanje stabla sa heš-mapom dece stablom sa dvostrukim nizom i uklanjanje režijskih troškova generičkog API-ja. Korišćenje postojeće bit mape po čvoru kao provere validnosti i podržavanje stabla ogromnim stranicama od 2 MB zajedno su obezbedili 10-20% uštede povrh veće promene na stablo sa dvostrukim nizom, u zavisnosti od veličine ulaza.

Gotovi Rust omotači oko biblioteka SentencePiece i IREE ilustruju istu stvar u obrnutom smeru. Povrh izvornih C i C++ tokenajzera, FFI granica i alokacije po pozivu koje omotači uvode vidljivi su u brojkama latencije. Izvorni kod u Rust-u za konačni Viterbijev enkoder dostupan je na adresi 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 pretprocesiranju

Gde referentna implementacija troši svoje vreme

Kako smo vršili merenja

Zašto su SentencePiece i IREE brži: Stablo sa dvostrukim nizom (Double-Array Trie)

Osnovna linija (Baseline): Port referentne implementacije bez alokacija

Optimizacija 1: Stablo sa dvostrukim nizom bez dodatnih apstraktnih slojeva

Optimizacija 2: Bitmapa i ugrađeno pakovanje

Optimizacija 3: Velike stranice za stablo

Završno poređenje

Završne misli

Reference

Citiranje