Unigrami tokenisaatori protsessori jõudluse parandamine

Me rakendasime oma Unigrami tokenisaatori täiesti nullist ümber, keskendudes jõudlusele. Tootmissisendi pikkuste juures vähendab uus kodeerija p50 latentsust umbes 5× võrreldes Hugging Face'i tokenizersi teegiga, ~2× võrreldes SentencePiece'iga (C++) ja ~1.5× võrreldes IREE tokenisaatoriga (C). Seejuures on staatilised kuhjaeraldused nullilähedased. Tootmises vähendas see meie järeldusahela CPU kasutusmahtu 5–6 korda ning kärpis ümberjärjestaja latentsusest kahekohalise arvu millisekundeid.

LLM-i järelduskulusid käsitletakse tavaliselt GPU-tööna: KV-vahemälud, tähelepanutuumad, ekspertide suunamine. Kuid märkimisväärne osa tootmisliikluses toimub palju väiksemate mudelitega, mis on piirimudelitest kaks kuni kolm suurusjärku väiksemad. Manustamismudelid, klassifitseerijad ja ümberjärjestajad töötavad endiselt GPU-del, kuid edasisuunamine (forward pass) on piisavalt lühike, et kõik selle ümber hakkab oluliseks muutuma.

Kõige selgem näide on ümberjärjestaja, mis hindab päringu kohta sadu kandidaatdokumente. Väikese mudeli puhul lõpeb GPU arvutustöö tavaliselt ühe-kohalise millisekundi jooksul, kuid iga sisend peab siiski esmalt läbima protsessoripoolse tokeniseerimise, mis muudab teksti sõnastiku ID-deks. Kui pakett on suur, võib see moodustada märkimisväärse osa päringu kogulatentsusest.

See töö on suunatud mudelile XLM-RoBERTa, millel on SentencePiece'iga treenitud 250 tuhande tokeni suurune Unigrami sõnastik. Peenhäälestatud RoBERTa perekonna kodeerijad on tavaline valik järjestus-, otsingu- ja sarnasusülesannete jaoks. Esitame siin ühelõimelise kodeerimislatentsuse; tootmises paralleelime tööd tuumade vahel, kasutades teeki rayon.

Kogu selle postituse vältel käsitleme võrdlusalusena Hugging Face'i teeki tokenizers crate, kuna see on vaikimisi Rusti tokenisaator, mida enamik meeskondi kasutab. Võrdluseks on lisatud ka SentencePiece (C++) ja IREE (C) kui kaks teist peamist tootmisrakendust.

Unigrami tokeniseerimine, mille Kudo tutvustas 2018. aastal (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) ja mis on kasutusel raamistikus SentencePiece, käsitleb segmenteerimist tõenäorima tee probleemina: igal sõnastiku tokenil on õpitud log-tõenäosus ning kodeerija valib segmenteerimise, mille tokenite skooride summa on suurim. Probleem lahendatakse Viterbi algoritmi (1967) abil, kus baidi positsioonid toimivad graafi kihtidena ja sõnastiku tokenid külgnevat baidivahemikku katvate servadena. Rekurents on ühemõõtmeline dünaamiline planeerimine (DP) baidi positsioonide kohal:

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

Välimine tsükkel on sisendi pikkuse suhtes lineaarne, mida igal positsioonil piirab pikim sõnastiku token (XLM-RoBERTa puhul 48 UTF-8 baiti). Sisemine puu läbikäik (trie walk) on kriitiline tee: 16k tokeni pikkuse sisendi korral teeb see sadu tuhandeid üleminekuid ja töötleb iga kodeerimise kohta kümneid miljoneid instruktsioone.

Hugging Face'i teek tokenizers v0.22 rakendab rekurentsi otse, kasutades tagurpidi otsinguks räsikaarti ja räsikaardi lapsi igas puu sõlmes:

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
}

Ja sisemine tsükkel (lihtsustatud kujul) – iga eesliiteotsingu vaste eraldab uue stringi ja teeb seejärel räsikaardi päringu:

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
}

Kood on loetav Rust, kuid lähtekoodi tasemel nähtamatult tehakse palju tööd. Järgmises jaotises analüüsitakse, kuhu see aeg kulub.

Enne Viterbi käivitamist läbib sisend Unicoode'i normaliseerimise (NFKC) ja sõnapiiride märgistamise. Profileerimine näitab, et meie stsenaariumides moodustab eeltöötlus alla 5% kodeerimise koguajast, mistõttu keskendume postituse ülejäänud osas Viterbi edasisuunamisele.

CPU koodi kiiremaks muutmiseks on kaks võimalust: täita vähem instruktsioone või saavutada kõrgem IPC (instruktsiooni tsükli kohta). Cachegrind mõõdab esimest, perf stat teist. Kõik selles postituses toodud optimeerimised keskenduvad ühele neist.

CPU-d seiskuvad neljal põhjusel: vahemälu tabamatused (L1d ~4 tsüklit, L3 tabamatus DRAM-i on sadu), sõltuvate laadimiste ahelad, mida järjestusväline mootor ei saa paralleelsemaks muuta, haru valesti ennustamised (~15–20 tsüklit tühistatakse vale oletuse kohta) ja mäluhalduri (allocator) kutsed. Võrdlusaluse unigrami mootori rakendus tabab kõiki nelja:

Pudelikael	Mehhanism	Mõõdetud mõju
Mälu eraldamine vaste kohta	String::from_utf8 vastendatud baitidel pluss AHashMapi otsing puu vaste kohta; tokeni ID on juba lehes olemas, kuid võrdlusalus ei salvesta seda seal	7295 eraldust 512 tokeni juures, 14 542 tokeni 1k juures, 299 171 tokeni 16k juures
Viida tagaajamine baiti kohta	AHashMap igas puu sõlmes, mis räsib järgmise baidi ja dereferentseerib sammu kohta kaks viita; Node struktuur pakib kokku ka külmad väljad, nagu is_leaf, tõmmates need igal külastusel vahemällu	L1d tabamatuse määr püsib madalal (~2,5% 1k juures), kuna mäluhalduri kuumad lehed elavad L1-s; kulu ilmneb sõltuva laadimise latentsusena
L2 vahemälu ummistumine pikkade sisendite korral	DP-tabel, eesliite puhver ja väljundtokenite loend eraldatakse iga kutse korral värskelt	L2 tabamatuse määr tõuseb 12%-lt 128 tokeni juures 50%-ni 16k juures, kuna kumulatiivne eraldatud maht ületab tuumaspetsiifilise 2 MB L2 vahemälu

Viitetokeniseerija kitsaskohtade, mehhanismide ja mõõdetud mõjude ülevaade.

Valdav kulu on mälu eraldamine iga vaste kohta. 1k tokeni juures sooritab võrdlusaluse rakendus umbes 7 miljonit instruktsiooni kodeerimise kohta IPC-ga 2,8, millest enamik kulub mäluhalduri kutsetele ja stringide teisendustele. L2 tabamatus on oluline alles pärast ~1k tokenit, kui kumulatiivne eraldusjälg ületab tuumaspetsiifilise vahemälu.

Viitetokeniseerija p50 ja p99 latentsus sisendsuuruse järgi.

Viitetokeniseerija L2-vahemälu möödalaskude määr sisendsuuruse järgi.

Baidid eraldisena ühe kodeerimise kohta viitetokeniseerijal erinevate sisendsuuruste lõikes.

Tokeni-põhine mälueraldus on konstantne: umbes 2 KB ja ~18 eraldust tokeni kohta, olenemata sisendi suurusest. Latentsuse mittelineaarsus tuleneb L2-st: kumulatiivsed baidid ületavad 2 MB umbes 1k tokeni juures, tabamatuse määr tõuseb 8%-lt umbes 50%-ni ja p99 kõver paindub sellele vastavalt.

Võrdleme kolme tootmistokenisaatorit, mis kõik käituvad sama Unigrami algoritmiga sama XLM-RoBERTa mudeli vastu:

Mootor	Keel	Puu
Hugging Face	Rust	HashMap-children eesliitepuu
SentencePiece	C++	Darts kahekordse massiivi puu
IREE tokenisaator	C	Darts kahekordse massiivi puu

Võrdluses kasutatud tokeniseerijate teostused, keeled ja trie-kujundused.

Seome SentencePiece'i ja IREE otseste C/C++ FFI-liideste kaudu, mitte nende Rusti teekide kaudu; teegid lisavad kutsekohtadele eraldisi ja stringide teisendusi, mis ei kuulu aluseks oleva teegi kulude hulka.

Kõik mootorid on ehitatud täielike optimeerimistega (O3, natiivne sihtimine, lingi aja optimeerimine). Mõõtmised tehti Intel Xeon Platinum 8488C protsessori ühel vaiksel tuumal. Iga mootori ja sisendi suuruse jaoks kogume seinakella latentsuse (p50 ja p99, 10 000 iteratsiooni pärast 1000 soojendust), Valgrind cachegrindi kaudu sooritatud instruktsioonid, IPC ja vahemälu tabamatuse määrad perf stat riistvaraloenduritest ning kuhjaeralduse arvud mõõdetud akna ümber oleva globaalse mäluhalduri liidese abil.

514 tokeni juures (lisandunud +2 512-le on BOS/EOS sisestus):

Rakendus	p50 latentsus	Instruktsioonid	IPC	Mälueraldused
Hugging Face	349 µs	3,60M	2,89	7295
SentencePiece	128 µs	1,83M	3,39	1559
IREE	112 µs	2,28M	3,06	1

Tokeniseerija latentsus, käskude arv, IPC ja eraldised 514 tokeni juures.

SP ja IREE saavutavad umbes kolmandiku võrdlusaluse latentsusest. Nad sooritavad umbes poole vähem instruktsioone, töötavad kõrgema IPC-ga ning eraldavad mälu ühe kuni kahe suurusjärgu võrra vähem.

SentencePiece ja IREE kasutavad mõlemad kahekordse massiivi puud, mille tutvustas Aoe 1989. aastal (Aoe, 1989; Karoonboonyanan). Puu on kodeeritud kahte lamedasse täisarvumassiivi, base ja check:

cpp

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

Alamotsing on:

cpp

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

Puu läbimise iga samm vajab ainult kahte massiivi laadimist, ühte täisarvude liitmist ja ühte võrdlust, ilma räsimise, viitade tagaajamise ja sõlme-põhise kuhjaeralduseta. Konstruktsioon paigutab iga sõlme nii, et selle lapsed satuvad valideerimismassiivi vabadele kohtadele, võimaldades mitmel vanemal piirkondi jagada seni, kuni nende lapsed erinevad järgneva baidi poolest.

XLM-RoBERTa 250K sõnastiku jaoks mahub kogu puu umbes 9 MB külgnevasse mällu (~780K pesa × 12 baiti nihke, valideerimise ja tokeni ID jaoks pesa kohta). Ühe kodeerimise kuum töökomplekt on suurusjärgus 100 KB, mahtudes mugavalt L2 vahemällu ning riistvaraline eellaadija lukustub korrapärasele juurdepääsumustrile.

Võrdlusalus säilitab see-eest iga sõlme kohta laste räsikaardi, mis maksab umbes neli sõltuvat laadimist baidisammu kohta. Räsida bait, laadida korv, järgida viita korvikirjele, järgida teist viita kuhjas olevale laps sõlmele ning seejärel vaadata järgmise baidi jaoks oma räsikaarti. Kuhjaobjektid hajuvad ja enamik laadimisi ei taba vahemälu.

SP ja IREE salvestavad mõlemad tokeni ID otse puusõlme. Võrdlusaluse puu rakendus tagastab ainult vastendatud baidid, mistõttu kodeerija peab iga vaste jaoks eraldama uue stringi, otsima selle tokeni ID taastamiseks üles kõrval olevast räsikaardist ja lugema skoori eraldi vektorist. SP tagastab ID puu salvestatud väärtusena ja loeb skoori lamedast ujukoma-massiivist:

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 kasutab ühe kutse raames kõigi vastete tagastamise asemel olekuga kursorirakendust, kuid kulude mudel on sama: tokeni ID-d ja skoorid on kättesaadavad lamedate massiivide indeksitena ilma eraldamise või räsimiseta puu läbimisel. IREE eraldab kodeerimise kohta ainult ühe väljundpuhvri; SP eraldab DP ajal lisaks ühe võresõlme vaste kohta, tehes 514 tokeni juures kokku umbes 1559 eraldust.

Enne puu muutmist eraldame algoritmi kulu rakenduse kulust: kui kiiresti töötab võrdlusalus, kui tarbetu töö on eemaldatud?

Baastase hoiab alles võrdlusaluse räsikaardipõhiste lastega puu, kuid teeb kaks muudatust. Kutsuja omanduses olev puhverstruktuur (scratch struct) sisaldab DP-tabelit, eeltöödeldud baite ja väljundtokenite ID-sid, nii et kuum tee ei tee pärast soojendust ühtegi mälueraldust. Tokeni ID ja skoor salvestatakse otse puusõlme, eemaldades vaste-põhise stringi teisenduse ja sekundaarse räsikaardi otsingu.

Sisemine tsükkel:

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

See on sama algoritm, sama puu kuju ja sama juurdepääsumustriga, lihtsalt ilma mälueralduste ja liigsete otsinguteta. Tulemused 514 tokeni juures:

Mootor	p50	Instruktsioonid	IPC	Mälueraldused
Hugging Face	326 µs	3,66M	2,86	7295
SentencePiece	127 µs	1,81M	3,55	1559
IREE	111 µs	1,58M	3,26	1
Baastase	156 µs	1,53M	2,62	0

Nulleraldiste baasjoon võrrelduna olemasolevate tokeniseerijate teostustega 514 tokeni juures.

Vaste-põhise stringi materialiseerimise ja sekundaarse räsikaardi otsingu eemaldamine vähendab instruktsioone võrreldes võrdlusalusega 2,4 korda. IPC langeb nendega koos (2,86 → 2,62), kuna räsikaardi viida tagaajamine, mis varem oli mäluhalduri töö taha peidetud, on nüüd peamine kulukas tegevus baidi kohta.

Viite- ja nulleraldiste baasjoone tokeniseerija p50 ja p99 latentsus sisendsuuruse järgi.

Viite- ja nulleraldiste baasjoone tokeniseerija L2-vahemälu möödalaskude määr sisendsuuruse järgi.

Baastase vähendab p50 latentsust võrreldes võrdlusalusega igas suuruses 2–3 korda ning selle L2 tabamatus jääb kuni 2k tokeni puhul 2–3% juurde (võrdlusaluse 8–48% vastu). Ülejäänud kulu on räsikaardi viida tagaajamine, mida käsitleb järgmine jaotis.

Räsikaardi puu on viida tagaajamine. SP ja IREE lahendavad selle kahekordse massiivi puu abil (vt eelmist jaotist) ja me laenamame sama andmestruktuuri. Erinevus seisneb selles, et SP ja IREE on üldotstarbelised teegid, mis toetavad voogedastussisendit, väljundnihe(-seid) ja mitut kodeerimisrežiimi; see paindlikkus maksab neile lisaraamatupidamist puu ümber (võre SP-s, kaheksajärguline konveier IREE-s). Meil pole midagi sellest vaja, seega lisame puu otse Viterbi tsüklisse ja viskame ülejäänu ära.

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

Baidi samm maksab kaks massiivi lugemist, täisarvude liitmise ja võrdluse. Kogu puu on umbes 9 MB, kuid ükski kodeerimine ei läbi sellest rohkem kui ~100 KB, seega püsib kuum andmestik L2-s.

Mootor	p50	p99	Instruktsioonid	L2 tabamatus %	L3 tabamatus %
Hugging Face	327 µs	3340 µs	3,66M	8,17%	14,66%
SentencePiece	127 µs	1140 µs	1,82M	4,74%	4,07%
IREE	111 µs	132 µs	1,61M	5,82%	3,87%
Baastase	155 µs	1170 µs	1,53M	2,59%	16,66%
+ Darts	68 µs	80 µs	1,10M	6,45%	11,42%

Tokeniseerija latentsus, instruktsioonid ja vahemälu möödalaskude määr pärast double-array trie kasutuselevõttu.

Seinakella aeg langeb baastasemega võrreldes 2,3 korda ja võrdlusalusega võrreldes 4,8 korda. Instruktsioonid langevad baastasemega võrreldes veel 28%. Võre ja konveieri eemaldamine võimaldab meil sooritada 1,10M instruktsiooni kodeerimise kohta (SP 1,82M ja IREE 1,61M vastu) sama puu kuju juures; vahe on raamatupidamises, mida nad hoiavad üldotstarbeliste API-de jaoks, mida meil pole vaja.

Optimeerimise 1 Darts-puu vajab kahte sõltuvat laadimist baidisammu kohta. Esmalt vanema baasnihke, seejärel kontrollkirje (check entry) kandidaadist lapspesas, et kinnitada, et vanem omab seda. Kontrollkirje kinnitab ainult omandiõigust. Samale küsimusele saab vastata sõlme-põhise bitikaardiga, mis salvestab, millistel baitidel on lapsed. Kahekordse massiivi puu ehituse tõttu on bitt baidile b vanema bitikaardis seatud täpselt siis, kui selle baidi kandidaatpesa kuulub vanemale. Bitikaart ja kontrollmassiiv kannavad sama teavet, nii et saame kontrollmassiivi käitusaegsest struktuurist täielikult eemaldada.

Bitikaardi otsing kompileeritakse üheks abitestiks 64-bitise sõna vastu, mis töötab umbes ühe tsükliga. Seotud POPCNT-instruktsioon, mis loendab registris olevad seatud bitid, on see, mida teised bitikaardiga indekseeritud puudokumendid kasutavad lapse nihke arvutamiseks otse bitikaardist. Meil pole seda vaja, sest Darts-konstruktsioon salvestab selle nihke juba sõlme baasväljale.

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

Selle käigus pakime neli sõlme-põhist välja (bitikaart, baas, tokeni ID ja skoor) ühte 64-baidisesse vahemälu ritta. Üksik sõlme külastus toob nüüd sisse kõik: skoori ja tokeni ID, mida vajame lõppsõlmedes, ning bitikaardi ja baasi järgmise sammu jaoks. Kontrollmassiivi kasutatakse ainult ehitamise ajal ja see eemaldatakse käitusaegsest paigutusest. Kuumas tees on indekseerimine kontrollimata, kuna puu ehitus ja tsükli struktuur tagavad juba piirid. Ruum on kompromiss: 64 baiti sõlme kohta × ~780K pesa korral kasvab puu ~9 MB-lt ~50 MB-ni.

Tulemused 514 tokeni juures:

Mootor	p50	p99	Instruktsioonid	L2 pöördumised	L2 tabamatus %	L3 tabamatus %
Hugging Face	327 µs	3340 µs	3,66M	15,0K	8,17%	14,66%
SentencePiece	127 µs	1140 µs	1,82M	7,5K	4,74%	4,07%
IREE	111 µs	132 µs	1,61M	4,7K	5,82%	3,87%
Baastase	155 µs	1170 µs	1,53M	6,1K	2,59%	16,66%
+ Darts	68 µs	80 µs	1,10M	4,6K	6,45%	11,42%
+ Bitikaart ja sisemine pakkimine	65 µs	77 µs	1,06M	1,8K	20,53%	15,01%

Tokeniseerija tulemused 514 tokeni juures pärast bitmapi ja inline-pakkimise lisamist.

Seinakella aeg langeb 514 tokeni juures süstemaatiliselt veel 4,5% ning kasvab koos sisendi pikkusega, ulatudes 4k ja 16k juures 6–7%-ni. Kontrollmassiivi eemaldamine muudab baidipõhise puusammu kahe vahemälu rea laadimise asemel üheks, mis on 2,5 korda vähem L2 pöördumisi kodeerimise kohta (1,8K tavalise Dartsi 4,6K vastu). Tabamatus hüppab 6%-lt 20%-le, kuid suurem osa sellest on nimetaja: vähem kui poole vähemate pöördumiste juures võtavad esinevad tabamused suurema osa.

Puu andmed on umbes 50 MB (780 000 sõlme, igaüks 64 baiti). Tavalises Linuxi protsessis toetatakse seda mälu 4 KB lehtedega, mistõttu kogu puu katab umbes 12 000 erinevat virtuaalset lehte, mis on palju rohkem kui CPU esimese taseme andmete Translation Lookaside Buffer mahutab (Sapphire Rapidsil on 96 kirjet). Iga Viterbi samm puutub kokku erineva sõlmega ja 512 tokeni pikkuse kodeerimise jooksul kulutame lehitustabeli kõndimisele umbes 9000 tsüklit, mis on umbes 3% kodeerimise eelarvest.

Lahenduseks on toetada puud 4 KB lehtede asemel 2 MB hiiglaslike lehtedega. Sama 50 MB katab nüüd 25 lehte, mahub täielikult Translation Lookaside Bufferisse ning lehitustabeli kataloogid ise lõpetavad puuga konkureerimise L2 vahemälu nimel.

Linuxil on kaks viisi mälu toetamiseks 2 MB lehtedega. Opportunistlik viis on läbipaistvad hiiglaslikud lehed (THP): teie programm palub kernelil uuendada mälupiirkonda 4 KB lehtedelt 2 MB lehtedele ning taustal olev kernelilõim proovib seda teha, kui selleks on aega ja külgnevat füüsilist mälu saadaval. Teine võimalus on eraldada puu otse eelnevalt reserveeritud hiiglaslike lehtede kogumist (pool), mille kernel eraldab alglaadimisel:

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

Kogum peab hostil olema (vm.nr_hugepages määratud alglaadimisel). Meie tootmisinstantsides on reserveeritud 10 561 hiiglaslikku lehte; puu tarbib neist 24. Pärast käsku mmap kinnitab /proc/<pid>/smaps piirkonda: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Võit on sisendi suuruste lõikes 3–12% seinakella ajast, mis on järjepidev ja proportsionaalne töökomplekti pikkusega:

Tokenid	Ilma hiiglaslike lehtedeta	Hiiglaslike lehtedega	Δ
130	15,7 µs	15,2 µs	−3,6%
514	65,4 µs	63,1 µs	−3,4%
1026	129,5 µs	125,3 µs	−3,3%
4098	773 µs	679 µs	−12,0%
16 386	3808 µs	3554 µs	−6,8%

Tokeniseerija latentsus suurte lehtedega ja ilma nendeta erinevate sisendsuuruste korral.

Suurim võit on 4k tokeni juures, kus töökomplekt on piisavalt suur, et lehitustabeli liiklus konkureeris aktiivselt puu andmetega L2 ribalaiuse nimel. Üle 4k kasu kahaneb, kuna L3 tabamatus domineerib ja lehitustabeli käigud moodustavad koguajast väiksema osa.

Tuues kolm optimeerimist kokku, on siin lõplik mootor, mida on mõõdetud iga teise meie testitud RoBERTa tokenisaatori vastu. Lisaks C++ natiivsetele SentencePiece ja C IREE binaarfailidele võrdlesime ka valmiskujul Rusti teeke, mida tootmiskasutaja võiks valida: sentencepiece ja iree-tokenizer. Teek sentencepiece mõõdab 1,6–1,9 korda suuremat latentsust kui C++ natiivbinaarfail igal sisendi pikkusel, kusjuures ülekoormus tuleneb sellest, et teek eraldab igal kodeerimisel värske tokeni tükkide loendi, millest igaüks hoiab oma koopiat tüki tekstist. Teek iree-tokenizer seab oma sisemise puhvri suuruse vastavaks iga sisendiga, mis hoiab ära pideva kodeerimis-põhise eralduse, mida natiivne IREE binaarfail maksab sõltumata sisendi suurusest, ja see on kõige olulisem lühikeste sisendite puhul, kus eraldus muidu domineeriks tegeliku tokeniseerimistöö üle. Tootmissisendi pikkuste juures püsikulu amortiseerub ja teegi enda kutsungipõhine lisakulu mälufunktsioonide kopeerimisest tagasi Rusti surub selle latentsuse üle natiivbinaarfaili oma.

p50 ja p99 latentsus ühe kodeerimise kohta üle tokeniseerija teostuste ja sisendsuuruste.

Lõplik mootor käitab sama Viterbi algoritmi üle sama 250 000 tokeni sõnastiku, ei eralda kuumal teel midagi ja toodab võrdlusalusega tokeni-täpse väljundi. Iga baidi samm on nüüd üksainus 64-baidine vahemälu rea laadimine, üks bititest ja üks täisarvude liitmine, mis on vähem kui räsikaardi päring, mäluhalduri kutse ja mitu erijuhtude haru baastaseme rakenduses. Instruktsioonid kodeerimise kohta langesid 3,66M-lt 1,04M-ni, mis on 3,5-kordne langus. Seinakella aeg vastab sellele suhtele lühikeste sisendite juures ja laieneb veelgi pikkade sisendite puhul, kus võrdlusaluse tokeni-põhised eraldused hakkavad L2 ja L3 vahemälu täitma.

Praktikas vähendas see töö protsessori kasutusmahtu meie järeldusahela eeltöötluskihis 5–6 korda ning kärpis ümberjärjestaja latentsusest kahekohalise arvu millisekundeid.

Suurem osa latentsuse vähenemisest tuli kergetest muudatustest, mis eemaldasid raisatud töö: kodeerimispõhiste eralduste eemaldamine, räsikaardi lastega puu asendamine kahekordse massiiviga ja üldotstarbelise API ülekulu eemaldamine. Olemasoleva sõlme-põhise bitikaardi kasutamine kehtivuskontrollina ja puu toetamine 2 MB hiiglaslike lehtedega moodustasid koos 10–20% lisaks suuremale üleminekule kahekordse massiivi puule, olenevalt sisendi suurusest.

Valmiskujul Rusti teegid SentencePiece'i ja IREE ümber illustreerivad sama punkti tagurpidi. C ja C++ natiivsete tokenisaatorite peal on FFI piir ja kutsungipõhised eraldused, mida teegid sisse toovad, latentsusnumbrites nähtavad. Lõpliku Viterbi kodeerija Rusti lähtekood on saadaval aadressil 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}
}

Motivatsioon

Viterbi algoritm

Võrdlusaluse rakendus

Märkus eeltöötluse kohta

Kuhu võrdlusaluse rakendus oma aja kulutab

Kuidas me mõõtsime

Miks SentencePiece ja IREE on kiiremad: kahekordse massiivi puu

Baastaseme rakendus: võrdlusaluse mälueraldusevaba teisendus

Optimeerimine 1: Kahekordse massiivi puu ilma täiendavate abstraktsioonideta

Optimeerimine 2: bitmap ja sisemine pakkimine

Optimeerimine 3: trie jaoks suured lehed

Lõplik võrdlus

Lõppsõnad

Viited

Viide