Uzlabot Unigram Tokenizer CPU veiktspēju

Mēs no jauna izstrādājām savu Unigram tokenizer no nulles kā mērķtiecīgu veiktspējas projektu. Pie produkcijas ievades garumiem jaunais kodētājs samazina p50 latentumu par aptuveni 5× salīdzinājumā ar Hugging Face tokenizers crate, ~2× salīdzinājumā ar SentencePiece (C++) un ~1.5× salīdzinājumā ar IREE tokenizer (C), ar nulles stacionārā stāvokļa kaudzes alokācijām. Produkcijā tas samazināja CPU izmantošanu mūsu secināšanas steks par 5-6× un samazināja rerankera latentumu par divciparu milisekundēm.

LLM secināšanas izmaksas parasti tiek definētas kā GPU darbs: KV kešatmiņas, uzmanības kodoli, ekspertu maršrutēšana. Taču nozīmīga produkcijas datu plūsmas daļa darbojas ar daudz mazākiem modeļiem, kas ir par divām līdz trim lieluma pakāpēm mazāki nekā robežlīniju transformatori. Iegulšanas modeļi, klasifikatori un pārrankotāji joprojām darbojas uz GPU, bet pārsūtīšanas solis ir pietiekami īss, lai viss ap to sāktu iegūt nozīmi.

Visuzskatāmākais piemērs ir pārrankotājs, kas novērtē simtiem kandidātu dokumentu uz vienu pieprasījumu. Ar nelielu modeli GPU skaitļošana parasti beidzas viena cipara milisekundēs, taču katrai ievadei joprojām vispirms jāiziet caur CPU puses tokenizāciju, pārvēršot tekstu vārdnīcas ID. Kad pakete ir liela, tas var veidot nozīmīgu daļu no kopējā pieprasījuma latentuma.

Šis darbs ir vērsts uz XLM-RoBERTa — 250K tokenu Unigram vārdnīcu, kas apmācīta ar SentencePiece. Smalki noregulēti RoBERTa saimes kodētāji ir izplatīta izvēle ranžēšanas, izgūšanas un līdzības uzdevumiem. Mēs visā tekstā ziņojam par vienpavediena latentumu uz vienu kodēšanu; produkcijā mēs paralelizējam starp kodoliem, izmantojot rayon.

Visā šajā ierakstā mēs uzskatām Hugging Face tokenizers crate par atsauci, pret kuru veicam salīdzinošo testēšanu, jo tas ir noklusējuma Rust tokenizer, kuru izmanto lielākā daļa komandu. SentencePiece (C++) un IREE (C) ir iekļauti salīdzināšanai kā divas citas galvenās produkcijas implementācijas.

Unigram tokenizācija, ko 2018. gadā ieviesa Kudo (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) un ieviesa SentencePiece, segmentāciju strukturē kā visiespējamākā ceļa problēmu: katram vārdnīcas tokenam ir iemācīta žurnāla varbūtība, un kodētājs izvēlas segmentāciju, kuras tokenu rādītāji summējas līdz augstākajai vērtībai. Problēma tiek atrisināta ar Viterbi algoritmu (1967), kur baitu pozīcijas ir grafiku slāņi un vārdnīcas toniki ir malas, kas aptver nepārtrauktu baitu diapazonu. Rekurence ir viendimensijas DP pāri baitu pozīcijām:

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

Ārējais cikls ir lineārs attiecībā pret ievades garumu, un katrā pozīcijā to ierobežo garākais vārdnīcas tokens (48 UTF-8 baiti XLM-RoBERTa gadījumā). Iekšējā trie staigāšana ir karstais ceļš: 16k tokenu ievadē tā veic simtiem tūkstošu pāreju un izbeidz desmitiem miljonu instrukciju uz vienu kodējumu.

Hugging Face tokenizers v0.22 tieši implementē rekursiju, izmantojot jaucējtabulu apgrieztajai meklēšanai un jaucējtabulas bērnus katrā trie mezglā:

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
}

Un iekšējais cikls (vienkāršots) – katra prefiksa meklēšanas sakritība alocē jaunu virkni un pēc tam pārbauda jaucējtabulu:

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
}

Kods ir salasāms Rust, taču tiek paveikts daudz darba, kas ir neredzams avota līmenī. Nākamajā sadaļā ir aprakstīts, kur pazūd laiks.

Pirms Viterbi palaišanas ievade iziet cauri Unicode normalizācijai (NFKC) un vārdu robežu marķēšanai. Profilēšana uzrāda, ka priekšapstrāde aizņem mazāk nekā 5% no kopējā kodēšanas laika mūsu scenārijos, tāpēc atlikušajā ieraksta daļā mēs koncentrējamies uz Viterbi pārsūtīšanas soli.

Ir divi veidi, kā paātrināt CPU kodu: izpildīt mazāk instrukciju vai iegūt augstāku IPC. Cachegrind mēra pirmo, perf stat mēra otro. Katra optimizācija šajā ierakstā attiecas uz vienu vai otru.

CPU apstājas četru lietu dēļ: kešatmiņas trūkumu (L1d ~4 cikli, L3 trūkums līdz DRAM ir simtiem), atkarīgas slodzes ķēžu, kuras ārpus secības dzinējs nevar paralelizēt, zaru nepareizas paredzēšanas (ap 15-20 cikliem tiek iztīrīti uz katru nepareizo minējumu) un alokatora izsaukumu. Atsauces unigram dzinēja implementācija skar visas četras:

Sastrēgumvietas	Mehānisms	Izmērītā ietekme
Alokācija uz sakritību	String::from_utf8 uz saskaņotajiem baitiem plus AHashMap meklēšana uz katru trie sakritību; tokena ID jau atrodas lapā, bet atsauce to tur neuzglabā	7295 alokācijas pie 512 tokeniem, 14 542 pie 1k, 299 171 pie 16k
Rādītāja dzīšanās uz baitu	AHashMap katrā trie mezglā, veicot nākamā baita jaucējkodu un atsaucot divus rādītājus uz soli; Node struktūra arī apvieno aukstos laukus, piemēram, is_leaf, ievelkot tos kešatmiņā katrā apmeklējumā	L1d trūkumu līmenis saglabājas zems (~2.5% pie 1k), jo alokatora karstās lapas atrodas L1; izmaksas parādās kā atkarīgās slodzes latentums
L2 trāpīšana garām garām ievadēm	DP tabula, prefiksa melnraksts un izvades tokenu saraksts tiek no jauna alocēti katrā izsaukumā	L2 trūkumu līmenis pieaug no 12% pie 128 tokeniem līdz 50% pie 16k, jo kumulatīvā alokācijas pēda pārsniedz katra kodola 2 MB L2

Atsauces tokenizatora šaurās vietas, mehānismi un izmērītā ietekme.

Dominējošās izmaksas ir alokācija uz sakritību. Pie 1k tokeniem atsauces implementācija izbeidz (pabeidz) aptuveni 7M instrukciju uz vienu kodējumu ar IPC 2.8, lielāko daļu no tām alokatora izsaukumos un virkņu konvertācijās. L2 trūkumi ir svarīgi tikai pēc ~1k tokeniem, kad kumulatīvā alokācijas pēda pārpilda katra kodola kešatmiņu.

Atsauces tokenizatora p50 un p99 latentums atkarībā no ievades izmēra.

Atsauces tokenizatora L2 kešatmiņas kļūdu līmenis atkarībā no ievades izmēra.

Baiti, kas piešķirti uz katru kodēšanu ar atsauces tokenizatoru dažādos ievades izmēros.

Alokācija uz vienu tokenu ir konstante: aptuveni 2 KB un ~18 alokācijas uz vienu tokenu neatkarīgi no ievades lieluma. Latentuma nonlinearitāte rodas no L2: kumulatīvie baiti šķērso 2 MB ap 1k tokeniem, trūkumu līmenis pieaug no 8% līdz ~50%, un p99 līkne izliecas, lai sakristu.

Mēs salīdzinām trīs produkcijas tokenizatorus, kuri visi izmanto to pašu Unigram algoritmu pret to pašu XLM-RoBERTa modeli:

Dzinējs	Valoda	Trie
Hugging Face	Rust	HashMap-children prefiksa trie
SentencePiece	C++	Darts dubultā masīva trie
IREE tokenizer	C	Darts dubultā masīva trie

Salīdzinājumā izmantotās tokenizatoru implementācijas, valodas un trie attēlojumi.

Mēs saistām SentencePiece un IREE, izmantojot tiešu C/C++ FFI, nevis to Rust ietvarus; ietvari pievieno izsaukuma alokācijas un virkņu konvertācijas, kas nav daļa no pamata bibliotēkas izmaksām.

Visi dzinēji ir būvēti ar pilnām optimizācijām (O3, vietējo mērķēšanu, saistīšanas laika optimizāciju). Mērījumi tiek veikti, piespraužot vienam klusam kodolam Intel Xeon Platinum 8488C. Katram dzinējam un ievades lielumam mēs apkopojam sienas pulksteņa latentumu (p50 un p99, 10 000 iterācijas pēc 1000 iesildīšanās), instrukcijas, kas izbeigtas, izmantojot Valgrind cachegrind, IPC un kešatmiņas trūkumu rādītājus no perf stat aparatūras skaitītājiem, un kaudzes alokācijas, ko saskaita globālais alokatora ietvars ap izmērīto logu.

Pie 514 tokeniem (+2 virs 512 ir BOS/EOS injekcija):

Implementācija	p50 latentums	Instrukcijas	IPC	Alokācijas
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

Tokenizatora latentums, instrukciju skaits, IPC un alokācijas pie 514 tokeniem.

SP un IREE sasniedz aptuveni trešdaļu no atsauces latentuma. Tie izbeidz aptuveni divas reizes mazāk instrukciju, darbojas ar augstāku IPC un alocē par vienu līdz divām lieluma pakāpēm mazāk.

SentencePiece un IREE abi izmanto dubultā masīva trie, ko 1989. gadā ieviesa Aoe (Aoe, 1989; Karoonboonyanan). Trie ir kodēts divos plakanos veselu skaitļu masīvos: base un check:

cpp

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

Bērna meklēšana ir šāda:

cpp

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

Katrai trie šķērsošanas darbībai ir nepieciešami tikai divi masīva nolasījumi, viena vesela skaitļa saskaitīšana un viena salīdzināšana, bez jaucējkoda, bez rādītāja dzīšanas un bez mezgla kaudzes alokācijas. Konstrukcija novieto katru mezglu tā, ka tā bērni iekrīt brīvos slotos validācijas masīvā, ļaujot vairākiem vecākiem koplietot reģionus, ja vien viņu bērni atšķiras ar baitu, kuram tie seko.

XLM-RoBERTa 250K vārdnīcai visa trie ietilpst ~9 MB nepārtrauktā atmiņā (~780K sloti × 12 baiti offset, validācijai un token ID uz slotu). Viena kodējuma karstās darba kopas lielums ir ap 100 KB, kas ērti atrodas L2, un aparatūras priekšielādētājs fiksē regulāro piekļuves modeli.

Turpretim atsauce uztur bērnu jaucējtabulu katram mezglam, kas maksā aptuveni četras atkarīgas slodzes uz katru baitu soli. Veic baitu jaucējkodu, ielādē kausu, seko rādītājam uz kausa ierakstu, seko citam rādītājam uz bērna mezglu kaudzē, pēc tam konsultējas ar savu jaucējtabulu par nākamo baitu. Kaudzes objekti izkliedējas, un vairums slodžu netrāpa kešatmiņā.

SP un IREE abi glabā tokena ID tieši trie mezglā. Atsauces trie implementācija atgriež tikai saskaņotos baitus, tāpēc kodētājam ir jāalocē jauna virkne katrai sakritībai, jāmeklē tā blakus jaucējtabulā, lai atgūtu tokena ID, un jālasa rādītājs no atsevišķa vektora. SP atgriež ID kā trie saglabāto vērtību un nolasa rādītāju no plakanā pludiņa masīva:

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 izmanto stāvokļa kursora API, nevis atgriež visas sakritības vienā izsaukumā, bet izmaksu modelis ir tāds pats: tokenu ID un rādītāji ir sasniedzami kā plakano masīvu indeksi, bez alokācijas vai jaucējkoda pašā trie pastaigā. IREE alocē tikai vienu izvades buferi uz kodējumu; SP papildus alocē režģa mezglu uz sakritību DP laikā, kopumā sasniedzot aptuveni 1559 alokācijas pie 514 tokeniem.

Pirms trie mainīšanas mēs izolējam algoritma izmaksas no implementācijas izmaksām: cik ātri darbojas atsauce ar noņemto nevajadzīgo darbu?

Bāzes līnija saglabā atsauces jaucējtabulas bērnu trie, bet veic divas izmaiņas. Zvanītājam piederoša melnraksta struktūra satur DP tabulu, iepriekš apstrādātus baitus un izvades tokenu ID, tāpēc karstais ceļš neveic alokācijas pēc iesildīšanās. Tokena ID un rādītājs tiek saglabāti tieši trie mezglā, noņemot virknes konvertāciju uz sakritību un sekundārās jaucējtabulas meklēšanu.

Iekšējais cikls:

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

Šis ir tas pats algoritms ar to pašu trie formu un to pašu piekļuves modeli, tikai bez alokācijām un liekām meklēšanām. Rezultāti pie 514 tokeniem:

Dzinējs	p50	Instrukcijas	IPC	Alokācijas
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
Bāzes līnija	156 µs	1.53M	2.62	0

Nulles alokāciju bāzlīnija salīdzinājumā ar esošajām tokenizatoru implementācijām pie 514 tokeniem.

Virknes materializācijas noņemšana uz sakritību un sekundārās jaucējtabulas meklēšana samazina instrukcijas par 2.4× salīdzinājumā ar atsauci. IPC samazinās kopā ar tiem (2.86 → 2.62), jo jaucējtabulas rādītāja dzīšanās, kas iepriekš bija paslēpta aiz alokatora darba, tagad ir dominējošās izmaksas uz baitu.

Atsauces un nulles alokāciju bāzlīnijas tokenizatora p50 un p99 latentums atkarībā no ievades izmēra.

Atsauces un nulles alokāciju bāzlīnijas tokenizatora L2 kešatmiņas kļūdu līmenis atkarībā no ievades izmēra.

Bāzes līnija samazina p50 latentumu par 2–3× salīdzinājumā ar atsauci jebkurā izmērā, un tās L2 trūkumu līmenis paliek tuvu 2–3% līdz 2k tokenu pret atsauces 8–48%. Atlikušās izmaksas ir HashMap rādītāja dzīšanās, ko risina nākamā sadaļa.

Jaucējtabulas trie ir rādītāja dzīšanās. SP un IREE to atrisina ar dubultā masīva trie (skatīt iepriekšējo sadaļu), un mēs aizņemamies to pašu datu struktūru. Atšķirība ir tāda, ka SP un IREE ir vispārējas nozīmes bibliotēkas, kas atbalsta straumēšanas ievadi, izvades nobīdes un vairākus kodēšanas režīmus; šī elastība maksā papildu grāmatvedību ap trie (režģis SP, astoņu posmu cauruļvads IREE). Mums nav nepieciešams nekas no tā, tāpēc mēs iekļaujam trie tieši Viterbi ciklā un atmetam pārējo.

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

Baita solis maksā divus masīva nolasījumus, vesela skaitļa saskaitīšanu un salīdzināšanu. Visa trie ir ap 9 MB, bet jebkurš atsevišķs kodējums izstaigā tikai ~100 KB no tā, tāpēc karstie dati paliek L2.

Dzinējs	p50	p99	Instrukcijas	L2 trūkums%	L3 trūkums%
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%
Bāzes līnija	155 µs	1,170 µs	1.53M	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	6.45%	11.42%

Tokenizatora latentums, instrukcijas un kešatmiņas kļūdu līmenis pēc dubultmasīva trie ieviešanas.

Sienas pulkstenis samazinās par 2.3× no bāzes līnijas un 4.8× no atsauces. Instrukcijas samazinās par vēl 28% no bāzes līnijas. Režģa un cauruļvada samazināšana ļauj mums izbeigt 1.10M instrukciju uz kodējumu pret SP 1.82M un IREE 1.61M pie tās pašas trie formas; atšķirība ir grāmatvedība, ko tie saglabā vispārējas nozīmes API, kas mums nav nepieciešami.

Darts trie no 1. optimizācijas prasa divas atkarīgas slodzes uz baitu soli. Vispirms vecāka bāzes nobīde, pēc tam check ieraksts kandidāta bērna slotā, lai apstiprinātu, ka vecākam tas pieder. Check ieraksts tikai apstiprina piederību. Uz to pašu jautājumu var atbildēt ar mezglam raksturīgu bitu karti, kas reģistrē, kuriem baitiem ir bērni. Saskaņā ar dubultā masīva trie konstrukciju, bits baitam b vecāka bitu kartē ir iestatīts tieši tad, kad šī baita kandidātu slots pieder vecākam. Bitu karte un check masīvs satur vienu un to pašu informāciju, tāpēc mēs varam pilnībā atmest check masīvu no izpildlaika struktūras.

Bitu kartes meklēšana tiek kompilēta par vienu bitu testu pret vienu 64 bitu vārdu, kas darbojas apmēram ciklā. Saistītā POPCNT instrukcija, kas skaita visus iestatītos bitus reģistrā, ir tas, ko citas ar bitu karti indeksētas trie konstrukcijas izmanto, lai tieši aprēķinātu bērna nobīdi no bitu kartes. Mums tas nav vajadzīgs, jo Darts konstrukcija jau glabā šo nobīdi mezgla base laukā.

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

Kamēr mēs to darām, mēs iepakojam četrus laukus uz mezglu (bitu karte, base, token ID un score) vienā 64 baitu kešatmiņas rindā. Viens mezgla apmeklējums tagad ievelk visu: score un token ID, kas mums nepieciešami gala mezglos, kā arī bitu karti un base nākamajam solim. Check masīvs tiek izmantots tikai konstrukcijas laikā un ir atmests no izpildlaika izkārtojuma. Karstā ceļa indeksēšana ir neatzīmēta, jo trie konstrukcija un cilpas struktūra jau garantē robežas. Kompromiss ir Telpa: ar 64 baitiem uz mezglu × ~780K slotiem trie aug no ~9 MB līdz ~50 MB.

Rezultāti pie 514 tokeniem:

Dzinējs	p50	p99	Instrukcijas	L2 piekļuves	L2 trūkums%	L3 trūkums%
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%
Bāzes līnija	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%
+ Bitu karte un iekļautā iepakošana	65 µs	77 µs	1.06M	1.8K	20.53%	15.01%

Tokenizatora rezultāti pie 514 tokeniem pēc bitkartes un iekļautās iesaiņošanas pievienošanas.

Sienas pulkstenis sistemātiski samazinās par vēl 4.5% pie 514 tokeniem un pieaug ar ievades garumu, sasniedzot 6-7% pie 4k un 16k. Check masīva nomešana pārvērš katra baita trie soli no divām kešatmiņas rindas slodzēm uz vienu, kas ir 2.5× mazāk L2 piekļuvju uz kodējumu (1.8K pret vienkāršo Darts 4.6K). Trūkumu līmenis lec no 6% līdz 20%, bet lielākā daļa no tā ir saucējs: ar mazāk nekā uz pusi mazāk piekļuvju trūkumi, kas mums ir, aizņem lielāku daļu.

Trie dati ir aptuveni 50 MB (780 000 mezglu pa 64 baitiem katrs). Noklusējuma Linux procesā šī atmiņa tiek nodrošināta ar 4 KB lapām, tāpēc pilna trie aptver aptuveni 12 000 atšķirīgu virtuālo lapu, kas ir tālu pagātnē tam, ko CPU pirmā līmeņa datu Translation Lookaside Buffer var turēt (Sapphire Rapids ir 96 ieraksti). Katrs Viterbi solis skar citu mezglu, un vairāk nekā 512 tokenu kodējumā mēs sadedzinām aptuveni 9000 ciklu lapu tabulas pastaigās, kas ir apmēram 3% no budžeta uz vienu kodējumu.

Risinājums ir nodrošināt trie ar 2 MB milzīgām lapām, nevis 4 KB lapām. Tas pats 50 MB tagad aptver 25 lapas, pilnībā ietilpst Translation Lookaside Buffer, un pašas lapu tabulas direktoriji pārtrauc konkurēt ar trie par L2 kešatmiņu.

Linux ir divi veidi, kā nodrošināt atmiņu ar 2 MB lapām. Oportūnistiskais veids ir caurspīdīgās milzīgās lapas (THP): jūsu programma lūdz kodolam jaunināt atmiņas reģionu no 4 KB uz 2 MB lapām, un fona kodola pavediens mēģina to izdarīt, kad tam ir laiks un pieejama blakus esošā fiziskā atmiņa. Otra alternatīva ir alocēt trie tieši no iepriekš rezervēta milzīgu lapu baseina, ko kodols atliek sāknēšanas laikā:

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

Baseinam ir jāpastāv resursdatorā (vm.nr_hugepages iestatīts sāknēšanas laikā). Mūsu produkcijas instancēs ir rezervētas 10 561 milzīgas lapas; trie patērē 24 no tām. Pēc mmap, /proc/<pid>/smaps apstiprina reģionu: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Ieguvums ir 3–12% sienas pulkstenis visos ievades garumos, kas ir konsekvents un proporcionāls darba kopas garumam:

Tokeni	Bez milzīgām lapām	Ar milzīgām lapām	Δ
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%

Tokenizatora latentums ar un bez lielajām lapām dažādos ievades izmēros.

Lielākais ieguvums ir pie 4k tokeniem, kur darba kopa ir pietiekami liela, lai lapu tabulas trafiks aktīvi konkurētu ar trie datiem par L2 joslas platumu. Pārsniedzot 4k, ieguvums sarūk, jo L3 trūkumi dominē un lapu tabulas pastaigas kļūst par mazāku kopējā laika daļu.

Apvienojot trīs optimizācijas, šeit ir galvenais dzinējs, kas mērīts pret katru citu RoBERTa tokenizatoru, ko mēs testējām. Līdzās vietējiem C++ SentencePiece un C IREE binārajiem failiem mēs arī testējām gatavās Rust ietvaru bibliotēkas, kuras izmantotu produkcijas lietotājs: sentencepiece un iree-tokenizer. Sentencepiece pakotne mēra par 1.6-1.9× augstāku latentumu nekā vietējais C++ binārais fails jebkurā ievades garumā, un pieskaitāmās izmaksas rodas no tā, ka ietvars alocē jaunu tokenu gabalu sarakstu, no kuriem katrs satur savu gabala teksta kopiju, katrā kodējumā. Iree-tokenizer pakotne izmēra savu iekšējo buferi katrai ievadei, kas novērš pastāvīgu alokāciju uz katru kodējumu, ko vietējais IREE binārais fails maksā neatkarīgi no ievades lieluma, un tas ir vissvarīgāk īsām ievadēm, kur alokācija pretējā gadījumā dominētu faktiskajā tokenizācijas darbā. Produkcijas ievades garumos pastāvīgās izmaksas amortizējas, un ietvara paša izsaukuma pieskaitāmās izmaksas, kopējot tokenus atpakaļ Rust, paceļ tā latentumu virs vietējā binārā faila latentuma.

p50 un p99 latentums uz katru kodēšanu dažādās tokenizatoru implementācijās un ievades izmēros.

Galīgais dzinējs palaiž to pašu Viterbi virs tās pašas 250 000 tokenu vārdnīcas, nealocē neko uz karstā ceļa un ražo tokeniem identisku izvadi pret atsauci. Katrs baitu solis tagad ir viena 64 baitu kešatmiņas rindas ielāde, viena bitu kartes pārbaude un viena vesela skaitļa saskaitīšana, kas ir mazāk nekā HashMap pārbaude, alokatora izsaukums un vairāki īpašu gadījumu zari bāzes līnijas implementācijā. Instrukcijas uz vienu kodējumu samazinājās no 3.66M uz 1.04M, kas ir 3.5× kritums. Sienas pulkstenis atbilst šai attiecībai īsās ievadēs un vēl vairāk palielinās garās ievadēs, kur atsauces tokenu alokācijas sāk pārpildīt L2 un L3.

Praksē šis darbs samazināja CPU izmantošanu mūsu secināšanas steka priekšapstrādes slānī par 5-6× un samazināja rerankera latentumu par divciparu milisekundēm.

Lielākā daļa latentuma samazinājuma nāca no vieglajiem izmaiņām, kas novērsa izšķērdēto darbu: atmetot alokācijas uz kodējumu, aizstājot HashMap bērnu trie ar dubulto masīvu un noņemot vispārējā API pieskaitāmās izmaksas. Esošās mezgla bitu kartes izmantošana kā derīguma pārbaude un trie nodrošināšana ar 2 MB milzīgām lapām kopā veidoja 10-20% papildus lielākajai pārejai uz dubultā masīva trie, atkarībā no ievades lieluma.

Gatavie Rust ietvari ap SentencePiece un IREE ilustrē to pašu punktu pretēji. Papildus vietējiem C un C++ tokenizatoriem FFI robeža un izsaukuma alokācijas, ko ietvari ievieš, ir redzamas latentuma skaitļos. Galīgā Viterbi kodētāja Rust avots ir pieejams vietnē github.com/perplexityai/pplx-garden.

Hugging Face tokenizers — https://github.com/huggingface/tokenizers

SentencePiece — https://github.com/google/sentencepiece

IREE tokenizer — https://github.com/iree-org/iree

Kudo, T. (2018). Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates. https://arxiv.org/abs/1804.10959

Viterbi, A. J. (1967). Error bounds for convolutional codes and an asymptotically optimum decoding algorithm. IEEE Transactions on Information Theory. https://en.wikipedia.org/wiki/Viterbi_algorithm

Aoe, J. (1989). An efficient digital search algorithm by using a double-array structure. IEEE Transactions on Software Engineering. https://ieeexplore.ieee.org/document/31365/

Karoonboonyanan, T. An Implementation of Double-Array Trie. https://linux.thai.net/~thep/datrie/datrie.html

POPCNT (x86) — https://www.felixcloutier.com/x86/popcnt

Huge pages — https://docs.kernel.org/admin-guide/mm/hugetlbpage.html

BibTeX

@online{perplexity2026ImprovingUnigramTokenizerCpuPerformance,
  author = {{Perplexity} and Vladimir Zaytsev and Bex Brooker and Daisy Lee and Roshan Shankar},
  title = {{Improving unigram tokenizer CPU performance}},
  date = {2026-05-20},
  url = {https://www.perplexity.ai/hub/blog/improving-unigram-tokenizer-cpu-performance}
}

Motivācija

Viterbi algoritms

Atsauces implementācija

Piezīme par priekšapstrādi

Kur atsauces implementācija tērē savu laiku

Kā mēs mērījām

Kāpēc SentencePiece un IREE ir ātrāki: Dubultā masīva trie

Bāzes līnija: nulles alokācijas atsauces ports

Optimizācija 1: dubultā masīva trie bez papildu abstrakcijām

Optimizācija 2: bitkartes un iekļautā blīvēšana

Optimizācija 3: milzīgas lapas triecienam

Galīgais salīdzinājums

Noslēguma domas

Atsauces

Citāts