„Unigram“ tokenizatoriaus CPU našumo gerinimas

Mes iš naujo, nuo nulio, įgyvendinome savo „Unigram“ tokenizatorių kaip tikslinį našumo projektą. Esant gamybinio įvedimo ilgiams, naujasis koduotojas sumažina p50 delsgą maždaug 5 kartus, lyginant su „Hugging Face“ tokenizatorių paketu, ~2 kartus, lyginant su „SentencePiece“ (C++), ir ~1,5 karto, lyginant su IREE tokenizatoriumi (C), o stacionariosios būsenos krūvos (heap) paskirstymų iš viso nėra. Gamyboje tai sumažino CPU naudojimą mūsų išvedimo steke 5–6 kartus ir sutrumpino perrinktuvo (reranker) delsą dviženkliu milisekundžių skaičiumi.

LLM išvedimo kaštai paprastai suprantami kaip GPU darbas: KV podėliai, dėmesio branduoliai, ekspertų nukreipimas. Tačiau reikšminga gamybinio srauto dalis veikia su daug mažesniais modeliais – dviem ar trimis dydžio eilėmis mažesniais už pažangiausius transformatorius. Įterpimo modeliai, klasifikatoriai ir perrinktuvai vis dar veikia GPU, tačiau tiesioginis praėjimas (forward pass) yra pakankamai trumpas, kad viskas, kas yra aplink jį, pradėtų turėti reikšmės.

Akivaizdžiausias atvejis yra perrinktuvas, vertinantis šimtus kandidatų dokumentų vienai užklausai. Naudojant nedidelį modelį, GPU skaičiavimai paprastai baigiasi per vienaženklį milisekundžių skaičių, tačiau kiekviena įvestis vis tiek turi pirmiausia pereiti per CPU pusės tokenizavimą, paverčiant tekstą žodyno ID. Kai paketas didelis, tai gali sudaryti reikšmingą bendros užklausos delsos dalį.

Šis darbas skirtas „XLM-RoBERTa“ – 250 tūkst. tokenų „Unigram“ žodynui, išmokytam naudojant „SentencePiece“. Modifikuotos „RoBERTa“ šeimos koduotojai yra dažnas pasirinkimas reitingavimo, paieškos ir panašumo užduotims. Visur pateikiame vieno gijos kodavimo delsą; gamyboje lygiagretiname per branduolius naudodami rayon.

Šiame įraše „Hugging Face“ tokenizers paketą laikome etalonu, su kuriuo lyginame rezultatus, nes tai yra numatytasis „Rust“ tokenizatorius, kurį renkasi dauguma komandų. „SentencePiece“ (C++) ir IREE (C) įtraukti palyginimui kaip kiti du pagrindiniai gamybiniai įgyvendinimai.

„Unigram“ tokenizavimas, Kudo pristatytas 2018 m. (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) ir įdiegtas SentencePiece, pateikia segmentavimą kaip labiausiai tikėtino kelio problemą: kiekvienas žodyno tokenas turi išmoktą log-tikimybę, o koduotojas surenka segmentavimą, kurio tokenų įverčių suma yra didžiausia. Ši problema sprendžiama naudojant Viterbi algoritmą (1967 m.), kur baitų pozicijos yra grafo sluoksniai, o žodyno tokenai – briaunos, apimančios gretutinę baitų sritį. Rekursija yra vienmatis dinaminis programavimas per baitų pozicijas:

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

Išorinis ciklas yra tiesinis įvesties ilgiui, kiekvienoje pozicijoje apribotas ilgiausio žodyno tokeno (48 UTF-8 baitai „XLM-RoBERTa“ atveju). Vidinis medis (trie) yra karštasis kelias: 16 tūkst. tokenų įvedimui jis atlieka šimtus tūkstančių perėjimų ir įvykdo dešimtis milijonų instrukcijų vienam kodavimui.

„Hugging Face“ tokenizers v0.22 tiesiogiai įgyvendina rekursiją, naudojant maišos lentelę (hash-map) atvirkštinei paieškai ir maišos lentelės vaikus kiekviename medžio mazge:

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
}

O vidinis ciklas (supaprastintas) – kiekviena prefikso paieškos atitiktis išskiria naują tekstinę eilutę ir patikrina maišos lentelę:

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
}

Kodas yra lengvai skaitomas „Rust“, tačiau atliekama daug darbų, kurie yra nematomi pirminio kodo lygyje. Kitame skyriuje nagrinėjama, kur dingsta laikas.

Prieš paleidžiant Viterbi, įvestis pereina Unikodo normalizavimą (NFKC) ir žodžių ribų žymėjimą. Profiliavimas rodo, kad mūsų scenarijuose išankstinis apdorojimas užima mažiau nei 5 % bendro kodavimo laiko, todėl likusioje įrašo dalyje susitelkiame į Viterbi tiesioginį praėjimą.

Yra du būdai pagreitinti CPU kodą: vykdyti mažiau instrukcijų arba pasiekti didesnį IPC. Cachegrind matuoja pirmąjį, „perf stat“ – antrąjį. Kiekviena šiame įraše pateikta optimizacija priklauso vienam arba kitam.

CPU sustoja dėl keturių dalykų: podėlio praleidimų (L1d ~4 ciklai, L3 praleidimas iki DRAM yra šimtai), priklausomų užkrovimo grandinių, kurių netvarkingo vykdymo variklis negali lygiagretinti, šakojimosi prognozių klaidų (~15–20 ciklų išvaloma dėl kiekvieno neteisingo spėjimo) ir paskirstytojo iškvietimų. Etaloninis unigramų variklio įgyvendinimas susiduria su visais keturiais:

Susiaurėjusi vieta	Mechanizmas	Išmatuotas poveikis
Paskirstymas vienai atitikčiai	„String::from_utf8“ su atitinkamais baitais bei „AHashMap“ paieška kiekvienai medžio atitikčiai; tokeno ID jau yra lape, tačiau etaloninis sprendimas jo ten nesaugo	7 295 paskirstymai esant 512 tokenų, 14 542 esant 1k, 299 171 esant 16k
Rodyklių vaikymasis vienam baitui	„AHashMap“ kiekviename medžio mazge, maišuojanti kitą baitą ir išskirianti du žymeklius per žingsnį; „Node“ struktūra taip pat sujungia šaltus laukus, tokius kaip „is_leaf“, įtraukdama juos į podėlį kiekvieno apsilankymo metu	L1d praleidimų dažnis išlieka žemas (~2,5 % esant 1k), nes paskirstytojo karštieji puslapiai gyvena L1; kaina pasireiškia kaip priklausomo užkrovimo delsa
L2 strigimas esant ilgoms įvestims	DP lentelė, prefikso juodraštis ir išvedamų tokenų sąrašas naujai paskirstomi kiekvieno iškvietimo metu	L2 praleidimų dažnis pakyla nuo 12 % esant 128 tokenams iki 50 % esant 16k, kai bendras paskirstymo dydis viršija vienam branduoliui skirtą 2 MB L2 podėlį

Nuoroda į žetonizatoriaus siaurąsias vietas, mechanizmus ir išmatuotą poveikį.

Didžiausia kaina yra paskirstymas vienai atitikčiai. Esant 1k tokenų, etaloninis įgyvendinimas įvykdo maždaug 7 mln. instrukcijų vienam kodavimui, o IPC yra 2,8; dauguma jų – paskirstytojo iškvietimuose ir tekstinių eilučių konversijose. L2 praleidimai svarbūs tik viršijus ~1k tokenų, kai bendras paskirstymo dydis viršija branduolio podėlį.

Referencinio žetonizatoriaus p50 ir p99 vėlinimas pagal įvesties dydį.

Referencinio žetonizatoriaus L2 talpyklos praleidimų rodiklis pagal įvesties dydį.

Baitų, paskirstomų kiekvienam kodavimui referenciniu žetonizatoriumi įvairiuose įvesties dydžiuose, skaičius.

Paskirstymas vienam tokenui yra pastovus: apie 2 KB ir ~18 paskirstymų vienam tokenui, neatsižvelgiant į įvesties dydį. Delsos nelinijiškumas atsiranda dėl L2: bendri baitai viršija 2 MB ties ~1k tokenų, praleidimų dažnis pakyla nuo 8 % iki ~50 %, o p99 kreivė atitinkamai išsikreipia.

Lyginame tris gamybinius tokenizatorius, kurie visi naudoja tą patį „Unigram“ algoritmą su tuo pačiu „XLM-RoBERTa“ modeliu:

Variklis	Kalba	Medis (Trie)
Hugging Face	Rust	„HashMap“ vaikų prefikso medis
SentencePiece	C++	„Darts“ dvigubo masyvo medis
IREE tokenizatorius	C	„Darts“ dvigubo masyvo medis

Palyginime naudotos žetonizatorių įgyvendinimo versijos, kalbos ir trie struktūrų atvaizdai.

„SentencePiece“ ir IREE sujungiame per tiesioginį C/C++ FFI, o ne per jų „Rust“ „wrapper“ paketus; pastarieji prideda iškvietimo paskirstymus ir tekstinių eilučių konversijas, kurios nėra pagrindinės bibliotekos kaštų dalis.

Visi varikliai sukompiliuoti su visiškomis optimizacijomis (O3, vietinis taikymas, kompiliavimo laiko optimizavimas). Matavimai atlikti pririšus procesą prie vieno tylaus branduolio „Intel Xeon Platinum 8488C“ procesoriuje. Kiekvienam varikliui ir įvesties dydžiui renkame realųjį laiką (p50 ir p99, 10 000 iteracijų po 1 000 apšilimo), „Valgrind cachegrind“ įvykdytas instrukcijas, IPC bei podėlio praleidimų dažnius iš „perf stat“ aparatinės įrangos skaitiklių, o krūvos paskirstymus skaičiuoja pasaulinis paskirstytojo „wrapper“ išmatuotame lange.

Esant 514 tokenų (+2 prie 512 yra BOS/EOS įterpimas):

Įgyvendinimas	p50 delsa	Instrukcijos	IPC	Paskirstymai
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

Žetonizatoriaus vėlinimas, instrukcijų skaičius, IPC ir paskirstymai esant 514 žetonų.

SP ir IREE delsos trukmė sudaro maždaug trečdalį etaloninio sprendimo delsos. Jie įvykdo maždaug perpus mažiau instrukcijų, veikia esant didesniam IPC ir paskirsto nuo vieno iki dviejų dydžio eilių mažiau.

„SentencePiece“ ir IREE naudoja dvigubo masyvo medį, kurį 1989 m. pristatė Aoe (Aoe, 1989; Karoonboonyanan). Medis užkoduotas dviejuose plokščiuose sveikųjų skaičių masyvuose – „base“ ir „check":

cpp

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

Vaikinio elemento paieška yra:

cpp

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

Kiekvienam medžio perėjimo žingsniui reikia tik dviejų masyvo nuskaitymų, vieno sveikojo skaičiaus sudėties ir vieno palyginimo, be jokio maišos (hashing), be rodyklių vaikymosi ir be mazgo paskirstymo krūvoje. Konstrukcija patalpina kiekvieną mazgą taip, kad jo vaikiniai elementai patektų į laisvus validavimo masyvo lizdus, leidžiant keliems tėviniams elementams dalytis sritimis, kol jų vaikiniai elementai skiriasi sekančiu baitu.

„XLM-RoBERTa“ 250 tūkst. žodynui visas medis telpa į ~9 MB gretutmiesčio atminties (~780 tūkst. lizdų × 12 baitų poslinkiui, validavimui ir tokeno ID viename lizde). Vieno kodavimo aktyvusis darbinis rinkinys yra apie 100 KB dydžio, patogiai telpa L2 podėlyje, o aparatinės įrangos prefetcher išlaiko reguliarų prieigos ciklą.

Etaloninis sprendimas, priešingai, kiekviename mazge išlaiko vaikų maišos lentelę, o tai kainuoja maždaug keturis priklausomus užkrovimus baito žingsniui. Susekamas baitas, užkraunamas blokas (bucket), sekama žodžio nuoroda į bloko įrašą, sekama kita nuoroda į vaikinį mazgą krūvoje, tada tikrinama jo paties maišos lentelė kitam baitui. Krūvos objektai išsibarsto, o dauguma užkrovimų nepatenka į podėlį.

SP ir IREE saugo tokeno ID tiesiogiai medžio mazge. Etaloninė medžio realizacija grąžina tik atitinkamus baitus, todėl koduotojas privalomai išskiria naują tekstinę eilutę kiekvienai atitikčiai, ieško jos papildomoje maišos lentelėje, kad atgautų tokeno ID, ir nuskaito įvertį iš atskiro vektoriaus. SP grąžina ID kaip saugomą medžio reikšmę ir nuskaito įvertį iš plokščio slankiojojo kablelio masyvo:

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 naudoja būsenos žymeklio (cursor) API vietoj visų atitikčių grąžinimo vienu iškvietimu, tačiau kaštų modelis išlieka toks pat: tokeno ID ir įverčiai pasiekiami kaip plokščio masyvo indeksai, neatliekant jokio paskirstymo ar maišos pačios medžio paieškos metu. IREE paskirsto tik vieną išvedimo buferį vienam kodavimui; SP papildomai paskirsto gardelės mazgą kiekvienai atitikčiai dinaminio programavimo metu, iš viso atlikdama maždaug 1 559 paskirstymus esant 514 tokenų.

Prieš keičiant medį, atskiriame algoritmų ir realizacijos kaštus: kaip greitai veikia etaloninis sprendimas pašalinus nereikalingą darbą?

Bazinė linija išlaiko etaloninio sprendimo maišos lentelės vaikų medį, tačiau padaro du pakeitimus. Kvietėjo valdoma juodraščio struktūra saugo DP lentelę, apdorotus baitus ir išvedamus tokeno ID, todėl karštasis kelias po apšilimo neatlieka jokių paskirstymų. Tokeno ID ir įvertis saugomi tiesiogiai medžio mazge, taip pašalinant eilutės konversiją vienai atitikčiai ir papildomos maišos lentelės paiešką.

Vidinis ciklas:

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

Tai yra tas pats algoritmas, su tuo pačiu medžio pavidalu ir ta pačia prieigos struktūra, tik be jokių paskirstymų ir perteklinių paieškų. Rezultatai esant 514 tokenų:

Variklis	p50	Instrukcijos	IPC	Paskirstymai
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
Bazinė linija	156 µs	1,53M	2,62	0

Nulinio paskirstymo bazinė versija, palyginta su esamomis žetonizatoriaus įgyvendinimo versijomis esant 514 žetonų.

Atitikties eilučių materializavimo ir antrinės maišos lentelės paieškos pašalinimas sumažina instrukcijų skaičių 2,4 karto, lyginant su etaloniniu sprendimu. Kartu su jais krinta ir IPC (2,86 → 2,62), nes maišos lentelės rodyklių vaikymasis, anksčiau pasislėpęs už atminties paskirstymo darbų, dabar tampa didžiausiais kaštais vienam baitui.

Referencinio ir nulinio paskirstymo bazinio žetonizatoriaus p50 ir p99 vėlinimas pagal įvesties dydį.

Referencinio ir nulinio paskirstymo bazinio žetonizatoriaus L2 talpyklos praleidimų rodiklis pagal įvesties dydį.

Bazinė linija sumažina p50 delsą 2–3 kartus, lyginant su etaloniniu sprendimu, nepriklausomai nuo dydžio, o L2 praleidimų dažnis išlieka arti 2–3 % iki 2k tokenų (palyginimui su etaloniniu 8–48 %). Likę kaštai yra „HashMap“ rodyklių vaikymasis, kurį išsprendžia kitas skyrius.

Maišos lentelės medis yra rodyklių vaikymasis. SP ir IREE tai išsprendžia naudodami dvigubo masyvo medį (žr. ankstesnį skyrių), ir mes pasiskoliname tą pačią duomenų strukturą. Skirtumas tas, kad SP ir IREE yra universalios bibliotekos, palaikančios srautinį įvedimą, išvedimo poslinkius ir kelis kodavimo režimus; tas lankstumas jiems kainuoja papildomą apskaitą aplink medį (gardelė SP, aštuonių pakopų konvejeris IREE). Mums to nereikia, todėl medį tiesiogiai įterpiame (inline) į Viterbi ciklą, o visa kita pašaliname.

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

Baito žingsnis kainuoja du masyvo nuskaitymus, sveikojo skaičiaus sudėtį ir palyginimą. Visas medis yra apie 9 MB, tačiau bet kuris vienas kodavimas praeina tik ~100 KB jo dalį, todėl karštieji duomenys lieka L2 podėlyje.

Variklis	p50	p99	Instrukcijos	L2 praleidimai %	L3 praleidimai %
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 %
Bazinė linija	155 µs	1 170 µs	1,53M	2,59 %	16,66 %
+ Darts	68 µs	80 µs	1,10M	6,45 %	11,42 %

Žetonizatoriaus vėlinimas, instrukcijos ir talpyklos praleidimų rodikliai pritaikius dvigubo masyvo trie.

Realiųjų procesoriaus laiko trukmė sumažėja 2,3 karto, lyginant su bazine linija, ir 4,8 karto, lyginant su etaloniniu sprendimu. Instrukcijų skaičius sumažėja dar 28 %, lyginant su bazine linija. Gardelės ir konvejerio atsisakymas leidžia įvykdyti 1,10 mln. instrukcijų vienam kodavimui, palyginti su SP 1,82 mln. ir IREE 1,61 mln., esant tai pačiai medžio struktūrai; skirtumas – tai apskaita, kurios reikalauja universalios API, kurių mums nereikia.

„Darts“ medis iš 1 optimizacijos reikalauja dviejų priklausomų užkrovimų baito žingsniui. Pirmiausia tėvinio elemento bazinis poslinkis, tuomet tikrinimo įrašas kandidato vaiko lizde, kad būtų patvirtinta, jog tėvinis elementas jį valdo. Tikrinimo įrašas tik patvirtina nuosavybę. Į šį klausimą galima atsakyti naudojant mazgo bitų žemėlapį (bitmap), kuriame užfiksuota, kurie baitai turi vaikinius elementus. Pagal dvigubo masyvo medžio konstrukciją, bitas b-tajam baitui tėvinio elemento bitų žemėlapyje nustatomas tada, kai to baito kandidato lizdas priklauso tėviniam elementui. Bitų žemėlapis ir tikrinimo masyvas pateikia tą pačią informaciją, todėl galime visiškai atsisakyti tikrinimo masyvo vykdymo metu.

Bitų žemėlapio paieška tampa viena bitų patikra prieš vieną 64 bitų žodį, kuri įvykdoma maždaug per ciklą. Susijusi POPCNT instrukcija, kuri suskaičiuoja visus nustatytus bitus registre, yra tai, ką kiti bitų žemėlapio indeksuoti medžiai naudoja tiesiogiai vaikiniams poslinkiams skaičiuoti iš bitų žemėlapio. Mums to nereikia, nes „Darts“ konstrukcija tą poslinkį jau saugo mazgo baziniame lauke.

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

Be to, supakuojame keturis laukus vienam mazgui (bitų žemėlapį, bazę, tokeno ID ir įvertį) į vieną 64 baitų podėlio eilutę. Vieno mazgo apsilankymas dabar gauna viską: įvertį ir tokeno ID, kurių reikia galiniuose mazgams, bei bitų žemėlapį ir bazę kitam žingsniui. Tikrinimo masyvas naudojamas tik kūrimo metu ir yra pašalinamas iš vykdymo išdėstymo. Karštojo kelio indeksavimas yra netikrinamas, nes medžio kūrimas ir ciklo struktūra jau garantuoja ribas. Kompromisas yra erdvė: esant 64 baitams vienam mazgui × ~780 tūkst. lizdų, medis išauga nuo ~9 MB iki ~50 MB.

Rezultatai esant 514 tokenų:

Variklis	p50	p99	Instrukcijos	L2 prieigos	L2 praleidimai %	L3 praleidimai %
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 %
Bazinė 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 %
+ Bitų žemėlapis ir tiesioginis pakavimas	65 µs	77 µs	1,06M	1,8K	20,53 %	15,01 %

Žetonizatoriaus rezultatai esant 514 žetonų, pridėjus bitmap ir įterptinį paketavimą.

Realiųjų procesoriaus laiko trukmė sistemingai sumažėja dar 4,5 % esant 514 tokenų ir didėja kartu su įvesties ilgiu, pasiekdama 6–7 % esant 4k ir 16k. Tikrinimo masyvo atsisakymas pakeičia medžio žingsnį baitui iš dviejų podėlio eilučių užkrovimų į vieną, o tai reiškia 2,5 karto mažiau L2 prieigų vienam kodavimui (1,8K palyginti su 4,6K paprasto „Darts“ atveju). Praleidimų dažnis pašoka nuo 6 % iki 20 %, tačiau didžioji dalis to yra vardiklis: turint mažiau nei pusę prieigų, pasitaikantys praleidimai užima didesnę dalį.

Medžio duomenys užima maždaug 50 MB (780 000 mazgų po 64 baitus). Numatytajame „Linux“ procese ši atmintis palaikoma 4 KB puslapių, todėl visas medis apima apie 12 000 skirtingų virtualių puslapių – gerokai daugiau, nei CPU pirmojo lygio duomenų vertimų podėlis (TLB) gali talpinti („Sapphire Rapids“ turi 96 įrašus). Kiekvienas Viterbi žingsnis pasiekia skirtingą mazgą, o per 512 tokenų kodavimą mes sunaudojame maždaug 9 000 ciklų puslapių lentelės vaikščiojimams – tai apie 3 % vieno kodavimo biudžeto.

Sprendimas yra palaikyti medį 2 MB milžiniškais puslapiais (huge pages) vietoj 4 KB. Tas pats 50 MB dabar apima 25 puslapius, pilnai telpa TLB, o pačių puslapių lentelių katalogai nustoja konkuruoti su medžiu dėl L2 podėlio.

„Linux“ turi du būdus palaikyti atmintį su 2 MB puslapiais. Opportunistinis būdas yra skaidrūs milžiniški puslapiai (THP): jūsų programa paprašo branduolio atnaujinti atminties sritį iš 4 KB į 2 MB puslapius, o foninė branduolio gija bando tai padaryti, kai turi laiko ir laisvos gretutiniais principais paskirstytos fizinės atminties. Kita alternatyva – paskirstyti medį tiesiogiai iš iš anksto rezervuoto milžiniškų puslapių baseino, kurį branduolys atideda sistemos paleidimo metu:

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

Baseinas turi egzistuoti pagrindiniame kompiuteryje (vm.nr_hugepages nustatyta įkrovos metu). Mūsų gamybinėse instancijose rezervuota 10 561 milžiniškas puslapis; medis sunaudoja 24 iš jų. Po mmap, /proc/<pid>/smaps patvirtina sritį: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Laimėjimas yra 3–12 % realiojo laiko per įvesties dydžius, nuoseklus ir proporcingas darbinio rinkinio ilgiui:

Tokenai	Be milžiniškų puslapių	Su milžiniškais puslapiais	Δ
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 %

Žetonizatoriaus vėlinimas su milžiniškais puslapiais ir be jų skirtinguose įvesties dydžiuose.

Didžiausias laimėjimas yra ties 4k tokenų, kur darbinis rinkinys yra pakankamai didelis, kad puslapių lentelės srautas aktyviai konkuruotų su medžio duomenimis dėl L2 pralaidumo. Virš 4k laimėjimas mažėja, nes L3 praleidimai tampa dominuojantys, o puslapių lentelės vaikščiojimai sudaro mažesnę bendro laiko dalį.

Sujungus šias tris optimizacijas kartu, pateikiamas galutinis variklis, išmatuotas lyginant su kiekvienu kitu mūsų išbandytu „RoBERTa“ tokenizatoriumi. Kartu su gimtaisiais C++ „SentencePiece“ ir C IREE dvejetainiais failais ištyrėme ir paruoštus „Rust“ „wrapper“ paketus, kuriuos rinktųsi gamybinis vartotojas – „sentencepiece“ ir „iree-tokenizer“. „sentencepiece“ paketas rodo 1.6–1.9 karto didesnę delsą nei gimtasis C++ dvejetainis failas esant bet kokiam įvesties dydžiui, o šios išlaidos kyla dėl to, kad „wrapper“ paskirsto naują tokenų dalių sąrašą, kurių kiekvienas saugo savo dalies tekstinę kopiją kiekvieno kodavimo metu. „iree-tokenizer“ paketas pritaiko savo vidinį buferį pagal kiekvieną įvestį, taip išvengdamas pastovaus paskirstymo vienam kodavimui, kurį gimtasis IREE dvejetainis failas moka neatsižvelgiant į įvesties dydį, o tai svarbiausia trumpoms įvestims, kur paskirstymas kitu atveju dominuotų faktiniame tokenizavimo darbe. Esant gamybinėms įvesties ilgims, pastovi kaina amortizuojasi, o „wrapper“ iškvietimo kaštai nukopijuojant tokenus atgal į „Rust“ pakelia jo delsą virš gimtojo dvejetainio failo.

p50 ir p99 vėlinimas vienam kodavimui tarp žetonizatoriaus įgyvendinimo versijų ir įvesties dydžių.

Galutinis variklis vykdo tą patį Viterbi su tuo pačiu 250 000 tokenų žodynu, karštajame kelyje nieko nepaskirsto ir sukuria tokenų tikslumo išvedimą, sutampantį su etaloniniu. Kiekvienas baito žingsnis dabar yra vienas 64 baitų podėlio eilutės užkrovimas, viena bitų žemėlapio patikra ir viena sveikojo skaičiaus sudėtis, sumažėjusi nuo „HashMap“ paieškos, paskirstytojo iškvietimo ir kelių specialių atvejų šakojimų bazinėje realizacijoje. Instrukcijos vienam kodavimui krito nuo 3,66 mln. iki 1,04 mln. – tai 3,5 karto sumažėjimas. Realiųjų procesoriaus laiko trukmė atitinka šį santykį trumpoms įvestims ir dar labiau išsiplečia ilgoms įvestims, kur etaloninio sprendimo paskirstymai vienam tokenui pradeda perpildyti L2 ir L3.

Praktikoje šis darbas sumažino CPU naudojimą mūsų išvedimo steko išankstinio apdorojimo sluoksnyje 5–6 kartus ir sutrumpino perrinktuvo delsą dviženkliu milisekundžių skaičiumi.

Didžioji dalis delsos sumažinimo atsirado dėl lengvų pakeitimų, pašalinusių švaistomą darbą: atsisakius paskirstymų vienam kodavimui, pakeitus „HashMap“ vaikų medį dvigubu masyvu ir pašalinus bendrosios API kaštus. Esamo mazgo bitų žemėlapio naudojimas kaip galiojimo patikros ir medžio palaikymas 2 MB milžiniškais puslapiais bendrai sudarė 10–20 % papildomai prie didesnio perėjimo prie dvigubo masyvo medžio, priklausomai nuo įvesties dydžio.

Paruošti „Rust“ „wrapper“ paketai aplink „SentencePiece“ ir IREE iliustruoja tą patį principą atvirkščiai. Virš gimtųjų C ir C++ tokenizatorių, FFI riba ir „wrapper“ įvedami paskirstymai vienam iškvietimui yra matomi delsos skaičiuose. „Rust“ šaltinis galutiniam Viterbi koduotojui pasiekiamas adresu 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}
}

Motyvacija

Viterbi algoritmas

Etaloninis įgyvendinimas

Pastaba apie išankstinį apdorojimą

Kur etaloninis įgyvendinimas eikvoja laiką

Kaip matavome

Kodėl „SentencePiece“ ir IREE yra greitesni: dvigubo masyvo medis

Bazinė linija: nulinio paskirstymo etaloninio sprendimo realizacija

1 optimizacija: dvigubo masyvo medis be papildomų abstrakcijų

2 optimizavimas: bitų žemėlapio ir vidinis pakavimas

3 optimizavimas: milžiniški puslapiai triui

Galutinis palyginimas

Baigiamosios mintys

Nuorodos

Citata