Millora del rendiment de la CPU del tokenizer unigram

Hem reimplementat el nostre tokenizer Unigram des de zero com un projecte centrat en el rendiment. A longituds d'entrada de producció, el nou codificador redueix la latència p50 aproximadament 5× en comparació amb la crate de tokenizers de Hugging Face, ~2× en comparació amb SentencePiece (C++) i ~1,5× en comparació amb el tokenizer d'IREE (C), amb zero assignacions de monticle en estat estacionari. En producció, això ha reduït la utilització de CPU en la nostra pila d'inferència entre 5 i 6 vegades i ha eliminat mil·lisegons de doble dígit de la latència del re-ranking.

El cost de la inferència amb LLM normalment s'estructura com a treball de GPU: memòries cau KV, nuclis d'atenció, encaminament d'experts. Però una porció significativa del tràfic de producció s'executa en models molt més petits, de dos a tres ordres de magnitud més petits que els transformadors de frontera. Els models d'incrustació, els classificadors i els re-rankers encara s'executen a les GPU, però el pas endavant és prou curt perquè tot el que l'envolta comenci a importar.

El cas més clar és un re-ranker que puntua centenars de documents candidats per sol·licitud. Amb un model petit, el càlcul de la GPU normalment acaba en mil·lisegons d'un sol dígit, però cada entrada encara ha de passar primer per la tokenització del costat de la CPU, convertint el text en IDs de vocabulari. Quan el lot és gran, això pot representar una fracció significativa de la latència total de la sol·licitud.

Aquest treball s'adreça a XLM-RoBERTa, un vocabulari Unigram de 250.000 tokens entrenat amb SentencePiece. Els codificadors de la família RoBERTa ajustats amb precisió són una opció comuna per a tasques de classificació, recuperació i similitud. Informem sobre la latència per codificació d'un sol fil; en producció paral·lelitzem entre nuclis utilitzant rayon.

Al llarg d'aquesta publicació tractem la tokenizers crate de Hugging Face com la referència amb la qual fem el benchmarking, ja que és el tokenizer predeterminat en Rust al qual recorren la majoria d'equips. SentencePiece (C++) i IREE (C) s'inclouen per a la comparació com les altres dues implementacions principals de producció.

La tokenització Unigram, introduïda per Kudo el 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) i inclosa a SentencePiece, planteja la segmentació com un problema de camí més probable: cada token de vocabulari té una probabilitat logarítmica apresa, i el codificador tria la segmentació les puntuacions de la qual sumen el valor més alt. El problema es resol amb l'algorisme de Viterbi (1967), amb posicions de bytes com a capes de grafs i tokens de vocabulari com a vores que abasten un rang de bytes contigus. La recurrència és una programació dinàmica unidimensional sobre les posicions de bytes:

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

El bucle exterior és lineal respecte a la longitud de l'entrada, limitat a cada posició pel token de vocabulari més llarg (48 bytes UTF-8 per a XLM-RoBERTa). El recorregut del trie interior és el camí crític: en una entrada de 16k tokens, executa centenars de milers de transicions i retira desenes de milions d'instruccions per codificació.

La versió v0.22 dels tokenizers de Hugging Face implementa la recurrència directament, amb una taula hash per a la cerca inversa i fills de taula hash a cada node del trie:

rust

vocab:        Vec<(String, f64)>           // indexed by token ID
token_to_ids: AHashMap<String, u32>        // reverse lookup
trie:         Trie<u8>                     // byte-level prefix trie

struct Node<Label> {
    is_leaf: bool,
    children: AHashMap<Label, Node<Label>>, // heap-allocated map per node
}

I el bucle interior (simplificat): cada coincidència de cerca de prefix assigna una cadena nova i després prova la taula hash:

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
}

El codi és un Rust llegible, però passa molta feina que és invisible a nivell de codi font. La secció següent desglossa on va a parar el temps.

Abans que s'executi Viterbi, l'entrada passa per la normalització Unicode (NFKC) i el marcatge de límits de paraula. El perfilat situa el preprocessament en menys del 5% del temps total de codificació en els nostres escenaris, de manera que ens centrem en el pas endavant de Viterbi per a la resta de la publicació.

Hi ha dues maneres d'accelerar el codi de la CPU: executar menys instruccions o aconseguir un IPC més alt. Cachegrind mesura el primer, perf stat mesura el segon. Cada optimització d'aquesta publicació recau en l'un o en l'altre.

Les CPU s'encallen en quatre coses: errors de memòria cau (L1d ~4 cicles, un error d'L3 a DRAM és de centenars), cadenes de càrrega dependent que el motor fora d'ordre no pot paral·lelitzar, errors de predicció de salts (~15-20 cicles eliminats per cada predicció incorrecta) i crides a l'assignador. La implementació del motor unigram de referència té els quatre problemes:

Coll d'ampolla	Mecanisme	Impacte mesurat
Assignació per coincidència	String::from_utf8 sobre els bytes coincidents més una cerca a AHashMap per coincidència de trie; l'ID del token ja és a la fulla, però la referència no l'emmagatzema allè	7.295 assignacions a 512 tokens, 14.542 a 1k, 299.171 a 16k
Cerca de punters per byte	Un AHashMap a cada node de trie, fent un hash del següent byte i desreferenciant dos punters per pas; l'estructura Node també agrupa camps freds com is_leaf, arrossegant-los a la memòria cau en cada visita	La taxa d'errors d'L1d es manté baixa (~2,5% a 1k) perquè les pàgines calentes de l'assignador viuen a L1; el cost es mostra com a latència de càrrega dependent
Batuda d'L2 en entrades llargues	La taula DP, el bloc de temporals de prefix i la llista de tokens de sortida s'assignen de nou a cada crida	La taxa d'errors d'L2 puja del 12% a 128 tokens fins al 50% a 16k a mesura que la petjada d'assignació acumulada supera els 2 MB d'L2 per nucli

Colls d'ampolla, mecanismes i impactes mesurats del tokenizer de referència.

El cost dominant és l'assignació per coincidència. A 1k tokens, la implementació de referència retira (completa) aproximadament 7 milions d'instruccions per codificació a un IPC de 2,8, la majoria en crides a l'assignador i conversions de cadenes. Els errors d'L2 només importen passats els ~1k tokens, un cop la petjada d'assignació acumulada desborda la memòria cau per nucli.

Latència p50 i p99 del tokenizer de referència segons la mida d'entrada.

Taxa de fallades de memòria cau L2 del tokenizer de referència segons la mida d'entrada.

Bytes assignats per cada encode del tokenizer de referència a través de diferents mides d'entrada.

L'assignació per token és constant: uns 2 KB i ~18 assignacions per token, independentment de la mida de l'entrada. La no linealitat de la latència prové d'L2: els bytes acumulats superen els 2 MB al voltant de 1k tokens, la taxa d'errors puja del 8% al ~50%, i la corba p99 es corba per coincidir-hi.

Comparem tres tokenizers de producció, tots executant el mateix algorisme Unigram contra el mateix model XLM-RoBERTa:

Motor	Llenguatge	Trie
Hugging Face	Rust	Trie de prefix amb fills HashMap
SentencePiece	C++	Trie de doble matriu Darts
Tokenizer d'IREE	C	Trie de doble matriu Darts

Implementacions del tokenizer, llenguatges i representacions de trie utilitzades en la comparació.

Enllacem SentencePiece i IREE mitjançant FFI directe de C/C++ en lloc dels seus wrappers de Rust; els wrappers afegeixen assignacions per crida i conversions de cadenes que no formen part del cost de la biblioteca subjacent.

Tots els motors estan compilats amb optimitzacions completes (O3, orientació nativa, optimització en temps d'enllaç). Les mesures s'executen fixades a un nucli silenciós en un Intel Xeon Platinum 8488C. Per a cada motor i mida d'entrada, recollim la latència de rellotge de paret (p50 i p99, 10.000 iteracions després d'1.000 d'escalfament), les instruccions retirades mitjançant Valgrind cachegrind, l'IPC i les taxes d'errors de memòria cau dels comptadors de maquinari de perf stat, i les assignacions de monticle comptades per un wrapper d'assignador global al voltant de la finestra mesurada.

A 514 tokens (el +2 sobre 512 és la injecció BOS/EOS):

Implementació	Latència p50	Instruccions	IPC	Assignacions
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

Latència del tokenizer, nombre d'instruccions, IPC i assignacions a 514 tokens.

SP i IREE es queden aproximadament en un terç de la latència de la referència. Retiren aproximadament la mitad d'instruccions, s'executen amb un IPC més alt i assignen entre un i dos ordres de magnitud menys.

SentencePiece i IREE utilitzen el trie de doble matriu introduït per Aoe el 1989 (Aoe, 1989; Karoonboonyanan). El trie està codificat en dues matrius d'enters planes, base i check:

cpp

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

Una cerca de fills és:

cpp

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

Cada pas del recorregut del trie només necessita dues lectures de matriu, una suma d'enters i una comparació, sense hashing, sense recerca de punters i sense assignació de monticle per node. La construcció col·loca cada node de manera que els seus fills caiguin en ranures lliures de la matriu de validació, permetent que diversos pares comparteixin regions sempre que els seus fills difereixin en el byte que segueixen.

Per al vocabulari de 250K d'XLM-RoBERTa, tot el trie cap en ~9 MB de memòria contigua (~780K ranures × 12 bytes per a l'desplaçament, la validació i l'ID de token per ranura). El conjunt de treball actiu d'una sola codificació és de l'ordre de 100 KB, còmodament a L2, i el pre-carregador de maquinari es fixa en el patró d'accés regular.

La referència, per contra, manté una taula hash de fills per node, costant aproximadament quatre càrregues dependents per pas de byte. S'aplica un hash al byte, es carrega el cubell, es segueix un punter a l'entrada del cubell, es segueix un altre punter al node fill al monticle, i després es consulta la seva pròpia taula hash per al següent byte. Els objectes del monticle s'escampen i la majoria de càrregues fallen a la memòria cau.

SP i IREE emmagatzemen l'ID de token directament al node del trie. La implementació de trie de referència només retorna els bytes coincidents, de manera que el codificador ha d'assignar una cadena nova per a cada coincidència, buscar-la en una taula hash lateral per recuperar l'ID de token i llegir la puntuació d'un vector separat. SP retorna l'ID com a valor emmagatzemat del trie i llegeix la puntuació d'una matriu plana de coma flotant:

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 utilitza una API de cursor amb estat en lloc de retornar totes les coincidències en una sola crida, però el model de costos és el mateix: els IDs de tokens i les puntuacions són accessibles com a índexs de matrius planes, sense assignació ni hashing en el propi recorregut del trie. IREE només assigna un buffer de sortida per codificació; SP a més assigna un node de xarxa per coincidència durant la programació dinàmica, totalitzant aproximadament 1.559 assignacions a 514 tokens.

Abans de canviar el trie, aïllem el cost de l'algorisme del cost de la implementació: a quina velocitat s'executa la referència un cop eliminat el treball innecessari?

La línia base manté el trie amb fills de taula hash de la referència però fa dos canvis. Una estructura de temporals propietat de qui crida conté la taula de programació dinàmica, els bytes preprocessats i els IDs de tokens de sortida, de manera que el camí crític no fa cap assignació després de l'escalfament. L'ID de token i la puntuació s'emmagatzemen directament al node del trie, eliminant la conversió de cadenes per coincidència i la cerca a la taula hash lateral.

Bucle interior:

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

Aquest és el mateix algorisme, amb la mateixa forma de trie i el mateix patró d'accés, només que sense assignacions ni cerques redundants. Resultats a 514 tokens:

Motor	p50	Instruccions	IPC	Assignacions
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
Línia base	156 µs	1,53M	2,62	0

La línia base sense assignacions comparada amb les implementacions existents del tokenizer a 514 tokens.

Eliminar la materialització de cadenes per coincidència i la cerca secundària a la taula hash redueix les instruccions 2,4× en comparació amb la referència. L'IPC baixa amb elles (2,86 → 2,62) perquè la cerca de punters a la taula hash, prèviament amagada d संग्राम per darrere del treball de l'assignador, és ara el cost dominant per byte.

Latència p50 i p99 del tokenizer de referència i de la línia base sense assignacions segons la mida d'entrada.

Taxa de fallades de memòria cau L2 del tokenizer de referència i de la línia base sense assignacions segons la mida d'entrada.

La línia base redueix la latència p50 entre 2 i 3 vegades en comparació amb la referència a totes les mides, i la seva taxa d'errors d'L2 es manté al voltant del 2-3% fins a 2k tokens enfront del 8-48% de la referència. El cost restant és la cerca de punters a HashMap, de la qual s'ocupa la secció següent.

El trie de taula hash és una cerca de punters. SP i IREE resolen això amb un trie de doble matriu (vegeu la secció anterior) i manllevem la mateixa estructura de dades. La diferència és que SP i IREE són biblioteques d'ús general que admeten entrada en streaming, desplaçaments de sortida i múltiples modes de codificació; aquesta flexibilitat els costa una comptabilitat addicional al voltant del trie (una xarxa a SP, un pipeline de vuit etapes a IREE). No necessitem res d'això, així que incrustem el trie directament al bucle de Viterbi i descartem la resta.

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

Un pas de byte costa dues lectures de matriu, una suma d'enters i una comparació. Tot el trie ocupa uns 9 MB, però qualsevol codificació individual només en recorre ~100 KB, de manera que les dades calentes es queden a L2.

Motor	p50	p99	Instruccions	% errors L2	% errors L3
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%
Línia base	155 µs	1.170 µs	1,53M	2,59%	16,66%
+ Darts	68 µs	80 µs	1,10M	6,45%	11,42%

Latència del tokenizer, instruccions i taxes de fallades de memòria cau després d'adoptar el trie d'array doble.

El rellotge de paret cau 2,3× des de la línia base i 4,8× des de la referència. Les instruccions cauen un altre 28% respecte a la línia base. Eliminar la xarxa i el pipeline ens permet retirar 1,10 milions d'instruccions per codificació enfront de l'1,82M de SP i l'1,61M d'IREE amb la mateixa forma de trie; la diferència és la comptabilitat que mantenen per a APIs d'ús general que no necessitem.

El trie Darts de l'optimització 1 necessita dues càrregues dependents per pas de byte. Primer l'desplaçament base del pare, després l'entrada check a la ranura fill candidata per confirmar que el pare la posseeix. L'entrada check només valida la propietat. La mateixa pregunta es pot respondre amb un mapa de bits per node que registri quins bytes tenen fills. Per la construcció del trie de doble matriu, el bit del byte b al mapa de bits del pare s'activa exactament quan el pare posseeix la ranura candidata per a aquest byte. El mapa de bits i la matriu check transporten la mateixa informació, de manera que podem eliminar completament la matriu check de l'estructura de temps d'execució.

La cerca al mapa de bits es compila en una sola prova de bit contra una paraula de 64 bits, la qual cosa s'executa aproximadament en un cicle. La instrucció POPCNT relacionada, que compta tots els bits actius en un registre, és el que utilitzen altres dissenys de trie indexats per mapa de bits per calcular un desplaçament de fill directament des del mapa de bits. No ho necessitem perquè la construcció Darts ja emmagatzema aquest desplaçament al camp base del node.

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

Aprofitant l'ocasió, empaquetem els quatre camps per node (mapa de bits, base, ID de token i puntuació) en una sola línia de memòria cau de 64 bytes. Una sola visita a un node ara carrega tot d'una: la puntuació i l'ID de token que necessitem als nodes terminals, més el mapa de bits i la base per al següent pas. La matriu check només s'utilitza durant la construcció i s'elimina del disseny de temps d'execució. La indexació del camí crític no està verificada perquè la construcció del trie i l'estructura del bucle ja garanteixen els límits. El compromís és l'espai: a 64 bytes per node × ~780K ranures, el trie creix d'uns 9 MB a uns 50 MB.

Resultats a 514 tokens:

Motor	p50	p99	Instruccions	Accessos L2	% errors L2	% errors L3
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%
Línia base	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%
+ Mapa de bits i empaquetat en línia	65 µs	77 µs	1,06M	1,8K	20,53%	15,01%

Resultats del tokenizer a 514 tokens després d'afegir bitmap i empaquetament en línia.

El rellotge de paret cau sistemàticament un altre 4,5% a 514 tokens i creix amb la longitud de l'entrada, arribant al 6-7% a 4k i 16k. Eliminar la matriu check converteix el pas de trie per byte de dues càrregues de línia de memòria cau a una, la qual cosa representa 2,5 vegades menys accessos a L2 per codificació (1,8K enfront dels 4,6K del Darts senzill). La taxa d'errors salta del 6% al 20%, però la majoria d'això es deu al denominador: amb menys de la mitad d'accessos, els errors que tenim prenen una quota més gran.

Les dades del trie ocupen aproximadament 50 MB (780.000 nodes a 64 bytes cadascun). En un procés de Linux predeterminat, aquesta memòria està respaldada per pàgines de 4 KB, de manera que tot el trie abasta unes 12.000 pàgines virtuals diferents, molt més enllà del que el buffer d'associació de traducció de dades de primer nivell de la CPU pot contenir (Sapphire Rapids té 96 entrades). Cada pas de Viterbi toca un node diferent, i al llarg d'una codificació de 512 tokens consumim aproximadament 9.000 cicles en recorreguts de taules de pàgines, al voltant del 3% del pressupost per codificació.

La solució és respaldar el trie amb pàgines grans (huge pages) de 2 MB en lloc de les de 4 KB. Els mateixos 50 MB ara abasten 25 pàgines, caben completament al Translation Lookaside Buffer, i els propis directoris de taules de pàgines deixen de competir amb el trie per la memòria cau L2.

Linux té dues maneres de respaldar memòria amb pàgines de 2 MB. L'oportunista són les pàgines grans transparents (THP): el teu programa demana al nucli que actualitzi una regió de memòria de pàgines de 4 KB a 2 MB, i un fil en segon pla del nucli intenta fer-ho quan té temps i memòria física contigua disponible. L'altra alternativa és assignar el trie directament des d'un pool de pàgines grans prèviament reservat que el nucli separa a l'arrencada:

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

El pool ha d'existir a l'amfitrió (vm.nr_hugepages configurat a l'arrencada). A les nostres instàncies de producció, es reserven 10.561 pàgines grans; el trie en consumeix 24. Després del mmap, /proc/<pid>/smaps confirma la regió: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

El guany és d'un 3–12% en el rellotge de paret a través de les mides d'entrada, consistent i proporcional a la longitud del conjunt de treball:

Tokens	Sense pàgines grans	Amb pàgines grans	Δ
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%

Latència del tokenizer amb i sense pàgines gegants a través de diferents mides d'entrada.

El benefici més gran es troba a 4k tokens, on el conjunt de treball és prou gran com perquè el tràfic de taules de pàgines estigués competint activament amb les dades del trie per l'amplada de banda d'L2. Més enllà de 4k, el guany s'encongeix perquè els errors d'L3 dominen i els recorreguts de taules de pàgines es converteixen en una fracció menor del temps total.

Ajuntant les tres optimitzacions, aquí teniu el motor final mesurat en comparació amb qualsevol altre tokenizer RoBERTa que haguem provat. Juntament amb els binaris natius en C++ de SentencePiece i en C d'IREE, també hem fet un benchmarking de les crates de wrapper de Rust estàndard a les quals recorreria un usuari de producció, sentencepiece i iree-tokenizer. La crate sentencepiece mesura una latència entre 1,6 i 1,9 vegades superior a la del binari natiu en C++ a totes les mides d'entrada, amb la sobrecàrrega provinent del fet que el wrapper assigna una llista nova de trossos de tokens, cadascun amb la seva pròpia còpia del text del tros, a cada codificació. La crate iree-tokenizer dimensiona el seu buffer intern per a cada entrada, la qual cosa evita una assignació constant per codificació que el binari natiu d'IREE paga independentment de la mida de l'entrada, i això importa més en entrades curtes on l'assignació altrament dominaria el treball real de tokenització. A longituds d'entrada de producció, el cost constant s'amortitza i la sobrecàrrega pròpia de cada crida del wrapper per copiar els tokens de nou a Rust eleva la seva latència per sobre de la del binari natiu.

Latència p50 i p99 per encode a través d'implementacions del tokenizer i mides d'entrada.

El motor final executa el mateix Viterbi sobre el mateix vocabulari de 250.000 tokens, no assigna res al camí crític i produeix una sortida exacta de tokens en comparació amb la referència. Cada pas de byte és ara una sola càrrega de línia de memòria cau de 64 bytes, una prova de mapa de bits i una suma d'enters, en lloc d'una prova de HashMap, una crida a l'assignador i diverses branques de casos especials a la implementació base. Les instruccions per codificació van baixar de 3,66M a 1,04M, una caiguda de 3,5×. El rellotge de paret coincideix amb aquesta proporció a entrades curtes i s'amplia encara més a entrades llargues, on les assignacions per token de la referència comencen a desbordar L2 i L3.

A la pràctica, aquest treball va reduir la utilització de la CPU a la capa de preprocessament de la nostra pila d'inferència entre 5 i 6 vegades i va eliminar mil·lisegons de doble dígit de la latència del re-ranker.

El gruix de la reducció de latència va provenir de canvis lleugers que van eliminar el treball inútil: eliminar les assignacions per codificació, substituir el trie de fills de HashMap per un de doble matriu i eliminar la sobrecàrrega de l'API genèrica. L'ús del mapa de bits per node existent com a comprovació de validesa i el suport del trie amb pàgines grans de 2 MB combinats van representar un 10-20% addicional per sobre del canvi més gran a un trie de doble matriu, depenent de la mida de l'entrada.

Les crates de wrapper de Rust estàndard al voltant de SentencePiece i IREE il·lustren el mateix punt a la inversa. A més dels tokenizers natius en C i C++, la frontera FFI i les assignacions per crida que introdueixen els wrappers són visibles en els números de latència. El codi font en Rust per al codificador Viterbi final està disponible a github.com/perplexityai/pplx-garden.

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

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

Tokenizer d'IREE — 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

Pàgines grans — 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}
}

Motivació

L'algorisme de Viterbi

Implementació de referència

Una nota sobre el preprocessament

On la implementació de referència gasta el seu temps

Com hem mesurat

Per què SentencePiece i IREE són més ràpids: el Trie de doble matriu

Línia base: una portabilitat sense assignacions de la referència

Optimització 1: Trie de doble matriu sense abstraccions addicionals

Optimització 2: mapa de bits i empaquetatge en línia

Optimització 3: pàgines grans per al trie

Comparació final

Reflexions finals

Referències

Citació