Unigram Tokenizer CPU Teljesítményének Javítása

Újraimplementáltuk az Unigram tokenizerünket a semmiből, egy kifejezetten a teljesítményre fókuszáló projektként. A tipikus éles bemeneti hosszaknál az új kódoló nagyjából 5×-ösére csökkenti az p50 latenciát a Hugging Face tokenizers crate-hez képest, ~2×-esére a SentencePiece (C++) és ~1,5×-esére az IREE tokenizere (C) képest, nulla permanens (steady-state) kupac-allokációval (heap allocation). Éles környezetben ez 5-6×-ra csökkentette a CPU-felhasználást a következtetési (inference) stackünkben, és kétjegyű ezredmásodperceket faragott le a rangsoroló (reranker) latenciájából.

Az LLM következtetési költséget általában GPU-munkaként szokás beállítani: KV-gyorsítótárak (cache), figyelmi kernelek (attention kernel), szakértő-útválasztás (expert routing). Az éles forgalom jelentős része azonban sokkal kisebb modelleken fut – két-három nagyságrenddel kisebbeken, mint a frontier transformer-ek. A beágyazási (embedding) modellek, osztályozók (classifier) és rangsorolók (reranker) továbbra is GPU-kon futnak, de az előremenő menet (forward pass) elég rövid ahhoz, hogy minden más, ami körülötte van, számítani kezdjen.

A legnyilvánvalóbb példa a rangsoroló, amely kérelmenként jelöltdokumentumok százait pontozza. Egy kis modellnél a GPU-számítás jellemzően egyjegyű ezredmásodperc alatt lefut, de minden bemenetnek át kell esnie először a CPU-oldali tokenizaláson, ami a szöveget szótár-azonosítókká (vocabulary ID) alakítja. Nagy kötegek (batch) esetén ez a teljes kérelem latenciájának jelentős hányadát teheti ki.

Ez a munka az XLM-RoBERTa-t veszi célba, amely egy SentencePiece-szel tanított, 250 ezer tokens Unigram szótárat használ. A finomhangolt RoBERTa-családú kódolók gyakori választások rangsorolási, visszakeresési (retrieval) és hasonlósági feladatokhoz. A mérések során mindenhol egy szállal futó, kódolásonkénti latenciát közlünk; éles környezetben a rayon segítségével magok között párhuzamosítunk.

Ebben a bejegyzésben a Hugging Face tokenizers crate-jét tekintjük referenciának, amelyhez viszonyítunk, mivel ez az alapértelmezett Rust tokenizer, amelyhez a legtöbb csapat nyúl. A SentencePiece (C++) és az IREE (C) összehasonlításként szerepel, mint a másik két jelentős éles implementáció.

Az Unigram tokenizálást Kudo mutatta be 2018-ban (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates), és a SentencePiece is ezt alkalmazza. Ez a szegmentálást a leghosszabb út (most-probable-path) problémájaként fogja fel: minden szótári tokenhez tartozik egy betanult log-valószínűség, és a kódoló azt a szegmentálást választja ki, amelynek token-pontszámai a legmagasabb értéket adják össze. A problémát a Viterbi-algoritmus (1967) oldja meg, ahol a bájtpozíciók a gráf rétegei, a szótári tokenek pedig az élek, amelyek egy folytonos bájttartományt hidalnak át. A rekurziónk egy egydimenziós dinamikus programozás (DP) a bájtpozíciókon:

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

A külső ciklus lineáris a bemeneti hosszhoz képest, amit minden pozícióban a leghosszabb szótári token határol be (48 UTF-8 bájt az XLM-RoBERTa esetében). A belső trie bejárás a forró útvonal (hot path): egy 16e tokenes bemeneten ez átmenetek százezreit futtatja le, és kódolásonként utasítások tízmillióit teljesíti.

A Hugging Face tokenizers v0.22 közvetlenül implementálja a rekurziót, egy hash-térképpel (hash-map) a fordított kereséshez, valamint hash-map gyermekekkel (hash-map-children) minden trie-csomóponton:

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
}

És a belső ciklus (egyszerűsítve) – minden prefix-keresési találat egy új sztringet allokál, majd lekérdezi a hash-térképet:

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
}

A kód olvasható Rust, de rengeteg olyan háttérmunka történik, amely a forráskód szintjén nem látszik. A következő szakasz bemutatja, hova tűnik az idő.

Mielőtt a Viterbi elindulna, a bemenet Unicode-normalizáláson (NFKC) és szóhatár-jelölésen esik át. A profilozás szerint a forgatókönyveinkben az előfeldolgozás a teljes kódolási idő kevesebb mint 5%-át teszi ki, így a bejegyzés hátralévő részében a Viterbi előremenő menetére fókuszálunk.

A CPU-kód felgyorsításának két módja van: kevesebb utasítást kell végrehajtani, vagy magasabb IPC-t (utasítás per ciklus) kell elérni. A Cachegrind az elsőt méri, a perf stat a másodikat. A bejegyzésben szereplő összes optimalizáció az egyikre vagy a másikra fókuszál.

A CPU-k négy dolog miatt akadnak el (stall): gyorsítótár-tévesztések (cache miss) (L1d ~4 ciklus, egy L3-tévesztés a DRAM-ból több száz), függő betöltési láncok (dependent-load chain), amelyeket a sorrenden kívüli (out-of-order) motor nem tud párhuzamosítani, elágazás-előrejelzési tévesztések (branch mispredict) (~15-20 ciklus törlésenként téves tipp esetén), valamint allokátori hívások. A referencia unigram motor implementációja mind a négybe beleütközik:

Szűk keresztmetszet	Mechanizmus	Mért hatás
Allokáció találatonként	String::from_utf8 az illesztett bájtokon plusz egy AHashMap-keresés trie-találatonként; a token azonosítója már benne van a levélben (leaf), de a referencia nem ott tárolja	7295 allokáció 512 tokennél, 14 542 1k-nál, 299 171 16k-nál
Mutatókövetés (pointer chase) bájtonként	Egy AHashMap minden trie-csomóponton, amely hash-eli a következő bájtot, és két mutatót dereferál lépésenként; a Node struktúra a ritkábban használt mezőket (cold field), például az is_leaf-et is egybevonja, így minden látogatáskor behúzza őket a cache-be	Az L1d tévesztési arány alacsony marad (~2,5% 1k-nál), mert az allokátor forró oldalai az L1-ben élnek; a költség függő betöltési latenciaként jelentkezik
L2 thrashing hosszú bemeneteknél	A DP-tábla, a prefix munkaterület (scratch) és a kimeneti tokenlista mind újonnan van allokálva minden híváskor	Az L2 tévesztési arány 12%-ról (128 token) 50%-ra emelkedik 16k-nál, ahogy a kumulatív allokációs lábnyom meghaladja a magonkénti 2 MB-os L2 méretet

Referencia a szűk keresztmetszetek, mechanizmusok és mért hatások a tokenizálásban.

A domináns költség a találatonkénti allokáció. 1k tokennél a referencia-implementáció kódolásonként megközelítőleg 7 millió utasítást hajt végre 2,8-as IPC mellett, amelyek nagy része allokátori hívásokban és sztringkonverziókban történik. Az L2-tévesztések csak ~1k token felett számítanak, amint a kumulatív allokációs lábnyom túlcsordul a magonkénti gyorsítótáron.

A referencia tokenizáló p50 és p99 késleltetése bemenetméret szerint.

A referencia tokenizáló L2 gyorsítótár-találati hibaaránya bemenetméret szerint.

A referencia tokenizáló által encode-onként lefoglalt bájtok a bemenetméretek függvényében.

A tokenenkénti allokáció állandó: körülbelül 2 KB és ~18 allokáció tokenenként, a bemenet méretétől függetlenül. A latencia nemlinearitása az L2-ből fakad: a kumulatív bájtok 2 MB fölé mennek 1k token körül, a tévesztési arány 8%-ról ~50%-ra nő, és a p99 görbe ennek megfelelően hajlik el.

Három éles tokenizert hasonlítunk össze, amelyek mindannyian ugyanazt az Unigram algoritmus futtatják ugyanazon az XLM-RoBERTa modellen:

Motor	Nyelv	Trie
Hugging Face	Rust	HashMap-children prefix trie
SentencePiece	C++	Darts double-array trie
IREE tokenizer	C	Darts double-array trie

A összehasonlításban használt tokenizáló-implementációk, nyelvek és trie-reprezentációk.

A SentencePiece-t és az IREE-t közvetlen C/C++ FFI-n keresztül kapcsoljuk be a Rust burkolók (wrapper) helyett; a burkolók olyan hívásonkénti allokációkat és sztringkonverziókat adnak hozzá, amelyek nem részei az alacsony szintű könyvtár költségének.

Minden motort teljes optimalizációval építünk (O3, natív célpont, link-time optimization). A mérések egy csendes magra rögzítve futnak egy Intel Xeon Platinum 8488C-n. Minden motorhoz és bemeneti mérethez összegyűjtjük a valós idejű (wall-clock) latenciát (p50 és p99, 10 000 iteráció 1000 bemelegítés után), a Valgrind cachegrind által lekérdezett utasításokat, az IPC-t és a cache-miss arányokat a perf stat hardveres számlálóiból, valamint a heap allokációkat, amelyeket egy globális allokátor wrapper számol a mért ablakon belül.

514 tokennél (az 512 feletti +2 a BOS/EOS injekció):

Implementáció	p50 latencia	Utasítások	IPC	Allokációk
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

Tokenizáló késleltetés, utasításszám, IPC és allokációk 514 tokennél.

Az SP és az IREE a referencia latenciájának megközelítőleg egyharmadánál tart. Nagyjából feleannyi utasítást hajtanak végre, magasabb IPC-vel futnak, és egy-két nagyságrenddel kevesebbet allokálnak.

A SentencePiece és az IREE egyaránt az Aoe által 1989-ben bevezetett duplatömbös trie-t használja (Aoe, 1989; Karoonboonyanan). A trie két lapos egész szám tömbbe van kódolva, a base-be és a check-be:

cpp

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

Egy gyermek keresése a következő:

cpp

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

A trie bejárásának minden lépése csupán két tömbolvasást, egy egész összeadást és egy összehasonlítást igényel, hash-elés, mutatókövetés és csomópontonkénti heap allokáció nélkül. Az építési folyamat úgy helyez el minden csomópontot, hogy annak gyermekei a validációs tömb szabad helyeire essenek, így több szülő osztozhat a tartományokon, amíg a gyermekek az általuk követett bájtokban különböznek.

Az XLM-RoBERTa 250K-s szókészletéhez a teljes trie belefér körülbelül 9 MB folytonos memóriába (körülbelül 780K bejegyzés × 12 bájt a bejegyzésenkénti eltolás, ellenőrzés és tokenazonosító számára). Egyetlen encode „forró” munkahalmaza nagyjából 100 KB, kényelmesen elfér az L2-ben, és a hardveres előbetöltő rááll a szabályos hozzáférési mintázatra.

A referencia ezzel szemben csomópontonkénti gyermek hash-térképet tart fenn, ami bájtlépésenként nagyjából négy függő betöltésbe kerül. Hash-elni a bájtot, betölteni a vödröt (bucket), követni egy mutatót a vödör-bejegyzéshez, követni egy másik mutatót a heapen lévő gyermek csomóponthoz, majd megnézni a saját hash-térképét a következő bájtra. A heap objektumok szétszóródnak, és a betöltések többsége téveszti a cache-t.

Az SP és az IREE is közvetlenül a trie-csomópontban tárolja a token azonosítóját. A referencia trie implementációja csak az illeszkedő bájtokat adja vissza, így a kódolónak minden találathoz külön sztringet kell allokálnia, meg kell keresnie egy oldalsó hash-térképben a token azonosítójának visszanyeréséhez, és be kell olvasnia a pontszámot egy külön vektorból. Az SP azonosítóként adja vissza a trie tárolt értékét, és a pontszámot egy lapos lebegőpontos tömbből olvassa:

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

Az IREE állapotalapú kurzor API-t használ ahelyett, hogy az összes találatot egyetlen hívásban adná vissza, de a költségmodell ugyanaz: a token azonosítók és pontszámok lapos tömb indexekként érhetők el, allokáció vagy hash-elés nélkül magán a trie bejárásán. Az IREE kódolásonként csak egy kimeneti puffert allokál; az SP emellett találatonként egy hálózati csomópontot (lattice node) is allokál a DP során, ami 514 tokennél nagyságrendileg 1559 allokációt eredményez.

Mielőtt megváltoztatnánk a trie-t, elkülönítjük az algoritmus költségét az implementáció költségétől: milyen gyorsan fut a referencia, ha eltávolítjuk a felesleges munkát?

Az alapvonal megtartja a referencia hash-map gyermekes trie-ját, but két változtatást hajt végre. Egy hívó által birtokolt munkaterület (scratch struct) tárolja a DP-táblát, az előfeldolgozott bájtokat és a kimeneti token azonosítókat, így a hot path a bemelegítés után nem végez allokációt. A token azonosító és a pontszám közvetlenül a trie-csomópontban van tárolva, így megszűnik a találatonkénti sztringkonverzió és a másodlagos hash-térkép keresés.

Belső ciklus:

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

Ez ugyanaz az algoritmus, ugyanolyan trie alakkal és hozzáférési mintával, csak allokációk és redundáns keresések nélkül. Eredmények 514 tokennél:

Engine	p50	Utasítások	IPC	Allokációk
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
Alapvonal	156 µs	1,53M	2,62	0

A nulla-allokációs alapvonal az 514 tokennél meglévő tokenizáló-implementációkhoz viszonyítva.

A találatonkénti sztring materializáció és a másodlagos hash-térkép keresés eltávolítása 2,4×-re csökkenti az utasításokat a referenciához képest. Az IPC ezzel együtt esik (2,86 → 2,62), mivel a hash-térkép mutatókövetése, amely korábban az allokátori munka mögött rejtőzött, mostantól a domináns bájtonkénti költség.

A referencia és a nulla-allokációs alapvonal tokenizáló p50 és p99 késleltetése bemenetméret szerint.

A referencia és a nulla-allokációs alapvonal tokenizáló L2 gyorsítótár-találati hibaaránya bemenetméret szerint.

Az alapvonal minden méretben 2-3×-ra csökkenti a p50 latenciát a referenciához képest, és az L2 tévesztési aránya 2k tokenig 2-3% körül marad a referencia 8-48%-ához képest. A fennmaradó költség a HashMap mutatókövetése, amivel a következő szakasz foglalkozik.

A hash-térkép trie egy mutatókövetés. Az SP és az IREE ezt egy duplatömbös trie-val oldja meg (lásd az előző szakaszt), és mi is kölcsönözzük ugyanezt az adatstruktúrát. A különbség az, hogy az SP és az IREE általános célú könyvtárak, amelyek támogatják a streaming bemenetet, a kimeneti eltolásokat és a többszörös kódolási módokat; ez a rugalmasság extra könyvelési költséget ró a trie körül (egy hálót az SP-ben, egy nyolcfázisú pipeline-t az IREE-ben). Nekünk ezekre semmi szükségünk, így beágyazzuk a trie-t közvetlenül a Viterbi-ciklusba, a többit pedig elhagyjuk.

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

Egy bájtlépés két tömbolvasásba, egy egész összeadásba és egy összehasonlításba kerül. Az egész trie körülbelül 9 MB, de bármelyik egyedi kódolás csak ~100 KB-ot jár be belőle, így a forró adatok az L2-ben maradnak.

Engine	p50	p99	Utasítások	L2 miss%	L3 miss%
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%
Alapvonal	155 µs	1170 µs	1,53M	2,59%	16,66%
+ Darts	68 µs	80 µs	1,10M	6,45%	11,42%

Tokenizáló késleltetés, utasítások és gyorsítótár-találati hibaarányok a double-array trie bevezetése után.

A valós idejű futási idő 2,3×-ra esik az alapvonalhoz és 4,8×-ra a referenciához képest. Az utasítások további 28%-kal csökkennek az alapvonalhoz képest. A háló és a pipeline elhagyása lehetővé teszi, hogy kódolásonként 1,10 millió utasítást teljesítsünk az SP 1,82M- és az IREE 1,61M-es értékével szemben, ugyanazon trie alak mellett; a rés az az általános célú API-knak fenntartott könyvelés, amelyre nekünk nincs szükségünk.

Az 1. optimalizációból származó Darts trie két függő betöltést igényel bájtlépésenként. Először a szülő base eltolását, majd a check bejegyzést a lehetséges gyermek slotnál, hogy megerősítse, a szülő birtokolja-e azt. A check bejegyzés csak a tulajdonjogot validálja. Ugyanez a kérdés megválaszolható egy csomópontonkénti bittérképpel, amely rögzíti, hogy mely bájtoknak vannak gyermekei. A duplatömbös trie konstrukciója révén a szülő bittérképében a b bájthoz tartozó bit pontosan akkor van beállítva, amikor az adott bájt lehetséges slotját a szülő birtokolja. A bittérkép és a check tömb ugyanazt az információt hordozza, így a check tömböt teljesen elhagyhatjuk a futásidejű struktúrából.

A bittérkép keresése egyetlen bit tesztre redukálódik egyetlen 64 bites szón szemben, ami nagyjából egy ciklus alatt fut le. A kapcsolódó POPCNT utasítás, amely megszámolja az összes beállított bitet egy regiszterben, az a funkció, amelyet más bittérkép-indexelt trie tervek használnak a gyermek eltolás közvetlen kiszámítására a bittérképből. Nekünk erre nincs szükségünk, mivel a Darts felépítés már tárolja ezt az eltolást a csomópont base mezőjében.

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

Ha már itt tartunk, a négy csomópontonkénti mezőt (bitmap, base, tokenazonosító és score) egyetlen 64 bájtos cache line-ba csomagoljuk. Egyetlen csomópont meglátogatása most mindent betölt: a terminális csomópontoknál szükséges score-t és tokenazonosítót, valamint a bitmapet és a base-t a következő lépéshez. Az ellenőrző tömböt csak az építés során használjuk, és eltávolítjuk a futásidejű elrendezésből. A hot path-beli indexelés ellenőrzés nélküli, mert a trie felépítése és a ciklus szerkezete eleve garantálja a határokat. Az ár a helyigény: 64 bájt csomópontonként × körülbelül 780K bejegyzés mellett a trie mérete körülbelül 9 MB-ról körülbelül 50 MB-ra nő.

Eredmények 514 tokennél:

Engine	p50	p99	Utasítások	L2 accesses	L2 miss%	L3 miss%
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%
Alapvonal	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%
+ Bittérkép és beágyazott csomagolás	65 µs	77 µs	1,06M	1,8K	20,53%	15,01%

Tokenizáló eredmények 514 tokennél a bitmap és az inline packelés hozzáadása után.

A valós idejű futási idő szisztematikusan újabb 4,5%-ot esik 514 tokennél, és a bemenet hosszával együtt nő, eléri a 6-7%-ot 4k és 16k méretnél. A check tömb elhagyása a bájtonkénti trie lépést két cache-sor betöltésből eggyé csökkenti, ami 2,5× kevesebb L2 hozzáférést jelent kódolásonként (1,8K a sima Darts 4,6K-jához képest). A tévesztési arány 6%-ról 20%-ra urik, de ennek nagy része a nevező: a kevesebb mint felére csökkent hozzáférések mellett a fennmaradó tévesztések nagyobb arányt képviselnek.

A trie adatai körülbelül 50 MB-ot tesznek ki (780 000 csomópont egyenként 64 bájttal). Egy alapértelmezett Linux folyamaton ezt a memóriát 4 KB-os oldalak támogatják, így a teljes trie mintegy 12 000 különböző virtuális oldalt ölel fel, ami messze meghaladja a CPU első szintű adatok Translation Lookaside Bufferének (TLB) kapacitását (a Sapphire Rapids 96 bejegyzéssel rendelkezik). Minden Viterbi lépés más csomópontot érint, és egy 512 tokenes kódolás során nagyságrendileg 9000 ciklust égetünk el oldal-tábla bejárásokban, ami a kódolásonkénti keret körülbelül 3%-a.

A megoldás az, hogy a trie-t 4 KB-os oldalak helyett 2 MB-os óriás oldalakkal (huge pages) támogassuk meg. Ugyanez az 50 MB most 25 oldalt ölel fel, teljesen elfér a Translation Lookaside Bufferben, és maguk az oldal-tábla könyvtárak megszűnnek versengeni a trie-val az L2 cache-ért.

A Linuxnak két módja van a memóriát 2 MB-os lapokkal hátterezni. Az eseti megoldás a transparent huge pages (THP): a programod megkéri a kernelt, hogy frissítsen egy memóriaterületet 4 KB-os lapokról 2 MB-os lapokra, és egy háttérben futó kernel-szál megpróbálja ezt elvégezni, amikor ideje és elegendő összefüggő fizikai memória áll rendelkezésre. A másik lehetőség, hogy a trie-t közvetlenül egy előre lefoglalt huge-page poolból foglaljuk, amelyet a kernel rendszerindításkor különít el:

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

A poolnak léteznie kell a hoston (vm.nr_hugepages beállítva indításkor). Az éles példányainkon 10 561 óriás oldal van lefoglalva; a trie ezek közül 24-et emészt fel. Az mmap után a /proc/<pid>/smaps megerősíti a régiót: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

A nyereség 3–12% körüli valós idejű futási időben a különböző bemeneti méreteknél, konzisztensen és arányosan a munkakészlet hosszával:

Tokenek	Óriás oldalak nélkül	Óriás oldalakkal	Δ
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%

Tokenizáló késleltetés nagy oldalakkal és azok nélkül a bemenetméretek függvényében.

A legnagyobb nyereség 4k tokennél jelentkezik, ahol a munkakészlet elég nagy ahhoz, hogy az oldal-tábla forgalom aktívan versengjen a trie adataival az L2 sávszélességért. 4k felett a nyereség csökken, mivel az L3 tévesztések dominálnak, és az oldal-tábla bejárások a teljes idő kisebb részévé válnak.

A három optimalizációt összevonva íme a végső motor, összehasonlítva az összes többi általunk tesztelt RoBERTa tokenizerrel. A natív C++ SentencePiece és C IREE binárisok mellett benchmarkoltuk a polcról levehető Rust wrapper crate-eket is, amelyekhez egy éles felhasználó nyúlna, a sentencepiece-t és az iree-tokenizer-t. A sentencepiece crate 1,6-1,9× magasabb latenciát mér, mint a natív C++ bináris minden bemeneti méretben; a többletterhelés abból adódik, hogy a wrapper minden kódoláskor allokál egy új token darab listát (token piece), amelyek mindegyike a darab szövegének saját másolatát hordozza. Az iree-tokenizer crate minden bemenethez méretezi a belső pufferét, ami elkerüli a folyamatos kódolásonkénti allokációt, amit a natív IREE bináris fizet a bemeneti mérettől függetlenül, és ez a rövidebb bemeneteknél számít a legtöbbet, ahol az allokáció különben dominálná a tényleges tokenizálási munkát. Az éles bemeneti hosszaknál az állandó költség amortizálódik, és a wrapper saját hívásonkénti többletterhelése a tokenek Rustba való másolásából a latenciáját a natív binárisé fölé emeli.

Az egyes encode-ok p50 és p99 késleltetése a tokenizáló-implementációk és bemenetméretek függvényében.

A végső motor ugyanezt a Viterbit futtatja ugyanazon a 250 000 tokenes szótáron, semmit sem allokál a hot path-en, és tokeneiben pontos kimenetet produkál a referenciához képest. Minden bájtlépés mostantól egyetlen 64 bájtos cache-sor betöltés, egy bittérkép teszt és egy egész összeadás, ami alacsonyabb a baseline implementáció HashMap vizsgálatánál, allokátori hívásánál és számos speciális eset ágánál. A kódolásonkénti utasítások száma 3,66 millióról 1,04 millióra esett, ami 3,5×-es csökkenés. A valós idejű futási idő rövid bemeneteknél követi ezt az arányt, hosszú bemeneteknél pedig tovább növekszik, ahol a referencia tokenenkénti allokációi túlcsordulni kezdenek az L2-n és az L3-on.

A gyakorlatban ez a munka 5-6×-ra csökkentette a CPU-felhasználást a következtetési stackünk előfeldolgozási rétegében, és kétjegyű ezredmásodperceket faragott le a rangsoroló latenciájából.

A latencia csökkenésének zöme olyan könnyű súlyú változtatásokból fakadt, amelyek eltávolították az elpazarolt munkát: elhagyták a kódolásonkénti allokációkat, a HashMap gyermekes trie-t duplatömbösre cserélték, és eltávolították az általános API-k többletterhelését. A meglévő csomópontonkénti bittérkép validációs ellenőrzésként való használata és a trie 2 MB-os óriás oldalakkal való támogatása együtt a duplatömbös trie-ra való áttérésen felül további 10-20%-os javulást eredményezett, a bemenet méretétől függően.

A SentencePiece és az IREE körüli, polcról levehető Rust wrapper crate-ek fordítva illusztrálják ugyanezt a pontot. A natív C és C++ tokenizeren felül a FFI határ és a wrapper által bevezetett hívásonkénti allokációk látszanak a latencia értékeken. A végső Viterbi kódoló Rust forráskódja elérhető itt: 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áció

A Viterbi-algoritmus

Referencia-implementáció

Megjegyzés az előfeldolgozásról

Hová megy el a referencia-implementáció ideje?

Hogyan mértünk

Miért gyorsabb a SentencePiece és az IREE: a duplatömbös trie (Double-Array Trie)

Alapvonal (Baseline): a referencia nulla allokációs portja

1. optimalizáció: duplatömbös trie extra absztrakciók nélkül

2. optimalizálás: Bitkép és beágyazott csomagolás

3. optimalizálás: Hatalmas lapok a triéhez

Végső összehasonlítás

Záró gondolatok

Hivatkozások

Idézet