Migliorare le prestazioni CPU del tokenizer unigram

Abbiamo reimplementato il nostro tokenizer Unigram da zero come progetto mirato alle prestazioni. Con lunghezze di input di produzione, il nuovo encoder riduce la latenza p50 di circa 5× rispetto al crate tokenizers di Hugging Face, ~2× rispetto a SentencePiece (C++) e ~1.5× rispetto al tokenizer di IREE (C), con zero allocazioni heap in stato stazionario. In produzione, questo ha ridotto l'utilizzo della CPU nel nostro stack di inferenza di 5-6× e ha ridotto la latenza del reranker di decine di millisecondi.

Il costo dell'inferenza LLM è solitamente inquadrato come lavoro della GPU: cache KV, kernel di attenzione, routing degli esperti. Tuttavia, una parte significativa del traffico di produzione viene eseguita su modelli molto più piccoli, da due a tre ordini di grandezza inferiori rispetto ai transformer di frontiera. I modelli di embedding, i classificatori e i reranker funzionano ancora su GPU, ma il forward pass è abbastanza breve da far sì che tutto ciò che lo circonda inizi a contare.

Il caso più chiaro è un reranker che valuta centinaia di documenti candidati per richiesta. Con un modello piccolo, il calcolo della GPU si completa solitamente in millisecondi a una sola cifra, ma ogni input deve ancora passare attraverso la tokenizzazione lato CPU, trasformando il testo in ID del vocabolario. Quando il batch è grande, ciò può rappresentare una frazione significativa della latenza totale della richiesta.

Questo lavoro ha come target XLM-RoBERTa, un vocabolario Unigram da 250K token addestrato con SentencePiece. Gli encoder della famiglia RoBERTa ottimizzati sono una scelta comune per attività di ranking, recupero e similarità. Riportiamo la latenza per encode single-threaded; in produzione parallelizziamo tra i core utilizzando rayon.

In tutto questo post trattiamo il tokenizers crate di Hugging Face come il riferimento rispetto al quale effettuiamo il benchmarking, poiché è il tokenizer Rust predefinito a cui la maggior parte dei team si rivolge. SentencePiece (C++) e IREE (C) sono inclusi per confronto come le altre due principali implementazioni di produzione.

La tokenizzazione Unigram, introdotta da Kudo nel 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) e distribuita in SentencePiece, inquadra la segmentazione come un problema di percorso più probabile: ogni token del vocabolario ha una log-probabilità appresa e l'encoder sceglie la segmentazione i cui punteggi dei token sommano al valore più alto. Il problema viene risolto con l'algoritmo di Viterbi (1967), con le posizioni dei byte come livelli del grafo e i token del vocabolario come archi che coprono un intervallo contiguo di byte. La ricorrenza è una DP unidimensionale sulle posizioni dei byte:

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

Il ciclo esterno è lineare rispetto alla lunghezza dell'input, limitato in ogni posizione dal token del vocabolario più lungo (48 byte UTF-8 per XLM-RoBERTa). Il trie walk interno è il percorso critico: su un input di 16k token, esegue centinaia di migliaia di transizioni e ritira decine di milioni di istruzioni per encode.

I tokenizers v0.22 di Hugging Face implementano la ricorrenza direttamente, con una hash-map per la ricerca inversa e hash-map-children in ogni nodo 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
}

E il ciclo interno (semplificato) – ogni corrispondenza di ricerca del prefisso alloca una nuova stringa e poi sonda l'hash-map:

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
}

Il codice è in Rust leggibile, ma avviene molto lavoro che è invisibile a livello di codice sorgente. La sezione successiva analizza dove va a finire il tempo.

Prima dell'esecuzione di Viterbi, l'input passa attraverso la normalizzazione Unicode (NFKC) e la marcatura dei confini di parola. Il profiling colloca la pre-elaborazione a meno del 5% del tempo totale di encode nei nostri scenari, quindi per il resto del post ci concentriamo sul forward pass di Viterbi.

Esistono due modi per velocizzare il codice CPU: eseguire meno istruzioni o ottenere un IPC più elevato. Cachegrind misura il primo, perf stat misura il secondo. Ogni ottimizzazione in questo post approda a uno o all'altro.

Le CPU vanno in stallo per quattro motivi: cache miss (L1d ~4 cicli, un L3-miss-to-DRAM è centinaia), catene di carichi dipendenti che il motore out-of-order non può parallelizzare, branch mispredict (~15-20 cicli svuotati per errore di previsione) e chiamate all'allocatore. L'implementazione del motore unigram di riferimento li colpisce tutti e quattro:

Collo di bottiglia	Meccanismo	Impatto misurato
Allocazione per corrispondenza	String::from_utf8 sui byte corrispondenti più una ricerca AHashMap per corrispondenza del trie; l'ID del token è già nella foglia, ma il riferimento non lo memorizza lì	7.295 allocazioni a 512 token, 14.542 a 1k, 299.171 a 16k
Caccia ai puntatori per byte	Un AHashMap in ogni nodo del trie, che esegue l'hashing del byte successivo e dereferenzia due puntatori per passaggio; la struct Node raggruppa anche campi freddi come is_leaf, trascinandoli nella cache a ogni visita	Il tasso di L1d miss rimane basso (~2,5% a 1k) perché le pagine calde dell'allocatore vivono in L1; il costo si manifesta come latenza del carico dipendente
Thrashing L2 su input lunghi	La tabella DP, lo scratch del prefisso e la lista dei token di output vengono tutti appena allocati a ogni chiamata	Il tasso di L2 miss sale dal 12% a 128 token al 50% a 16k man mano che l'impronta di allocazione cumulativa supera l'L2 da 2 MB per core

Riferimento ai colli di bottiglia del tokenizzatore, ai meccanismi e agli impatti misurati.

Il costo dominante è l'allocazione per corrispondenza. A 1k token, l'implementazione di riferimento ritira (completa) circa 7M di istruzioni per encode a un IPC di 2,8, la maggior parte delle quali in chiamate all'allocatore e conversioni di stringhe. Gli L2 miss contano solo oltre ~1k token, una volta che l'impronta di allocazione cumulativa supera la cache per core.

Latenza p50 e p99 del tokenizzatore di riferimento in base alla dimensione dell'input.

Tasso di mancati colpi nella cache L2 del tokenizzatore di riferimento in base alla dimensione dell'input.

Byte allocati per ogni encode dal tokenizzatore di riferimento su diverse dimensioni di input.

L'allocazione per token è costante: circa 2 KB e ~18 allocazioni per token, indipendentemente dalla dimensione dell'input. La non linearità della latenza deriva dall'L2: i byte cumulativi superano i 2 MB intorno a 1k token, il tasso di miss sale dall'8% a ~50% e la curva p99 si piega di conseguenza.

Confrontiamo tre tokenizer di produzione, tutti che eseguono lo stesso algoritmo Unigram sullo stesso modello XLM-RoBERTa:

Motore	Linguaggio	Trie
Hugging Face	Rust	Trie di prefisso HashMap-children
SentencePiece	C++	Darts double-array trie
Tokenizer IREE	C	Darts double-array trie

Implementazioni del tokenizzatore, linguaggi e rappresentazioni trie utilizzati nel confronto.

Colleghiamo SentencePiece e IREE tramite FFI C/C++ diretto anziché i loro wrapper Rust; i wrapper aggiungono allocazioni per chiamata e conversioni di stringhe che non fanno parte del costo della libreria sottostante.

Tutti i motori sono costruiti con ottimizzazioni complete (O3, targeting nativo, link-time optimization). Le misurazioni vengono eseguite su un singolo core tranquillo di un Intel Xeon Platinum 8488C. Per ogni motore e dimensione dell'input raccogliamo la latenza wall-clock (p50 e p99, 10.000 iterazioni dopo 1.000 di warmup), istruzioni ritirate tramite Valgrind cachegrind, IPC e tassi di cache miss da contatori hardware perf stat e allocazioni heap contate da un wrapper dell'allocatore globale attorno alla finestra misurata.

A 514 token (i +2 rispetto a 512 sono l'iniezione BOS/EOS):

Implementazione	Latenza p50	Istruzioni	IPC	Allocazioni
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

Latenza del tokenizzatore, conteggio delle istruzioni, IPC e allocazioni a 514 token.

SP e IREE arrivano a circa un terzo della latenza del riferimento. Ritirano circa la metà delle istruzioni, funzionano a un IPC più elevato e allocano da uno a due ordini di grandezza in meno.

SentencePiece e IREE usano entrambi il double-array trie introdotto da Aoe nel 1989 (Aoe, 1989; Karoonboonyanan). Il trie è codificato in due array di interi piatti, base e 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 ricerca del figlio è:

cpp

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

Ogni passaggio della traversata del trie richiede solo due caricamenti di array, una somma di interi e un confronto, senza hashing, senza caccia ai puntatori e senza allocazione heap per nodo. La costruzione posiziona ogni nodo in modo che i suoi figli cadano in slot liberi nell'array di validazione, consentendo a più genitori di condividere regioni finché i loro figli differiscono sul byte che seguono.

Per il vocabolario da 250K di XLM-RoBERTa, l'intero trie si adatta a ~9 MB di memoria contigua (~780K slot × 12 byte per l'offset, la validazione e l'ID del token per slot). Il set di lavoro a caldo di un singolo encode è nell'ordine dei 100 KB, comodamente in L2, e l'hardware prefetcher si aggancia al pattern di accesso regolare.

Il riferimento, al contrario, mantiene una hash-map di figli per nodo, costando circa quattro carichi dipendenti per passaggio di byte. Esegue l'hashing del byte, carica il bucket, segue un puntatore all'entrata del bucket, segue un altro puntatore al nodo figlio sull'heap, quindi consulta la sua hash-map per il byte successivo. Gli oggetti nell'heap si disperdono e la maggior parte dei carichi manca la cache.

SP e IREE memorizzano entrambi l'ID del token direttamente nel nodo del trie. L'implementazione del trie di riferimento restituisce solo i byte corrispondenti, quindi l'encoder deve allocare una nuova stringa per ogni corrispondenza, cercarla in un'hash-map laterale per recuperare l'ID del token e leggere il punteggio da un vettore separato. SP restituisce l'ID come valore memorizzato del trie e legge il punteggio da un array di float piatto:

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 utilizza un'API a cursore stateful invece di restituire tutte le corrispondenze in una singola chiamata, ma il modello di costo è lo stesso: gli ID dei token e i punteggi sono raggiungibili come indici di array piatti, senza allocazione o hashing sulla camminata del trie stessa. IREE alloca solo un buffer di output per encode; SP alloca inoltre un nodo del reticolo per corrispondenza durante la DP, per un totale di circa 1.559 allocazioni a 514 token.

Prima di modificare il trie, isoliamo il costo dell'algoritmo dal costo dell'implementazione: quanto velocemente funziona il riferimento con il lavoro non necessario rimosso?

La baseline mantiene il trie hash-map-children del riferimento ma apporta due modifiche. Una struct di scratch posseduta dal chiamante contiene la tabella DP, i byte pre-elaborati e gli ID dei token di output, quindi il percorso critico non effettua allocazioni dopo il warmup. L'ID del token e il punteggio sono memorizzati direttamente nel nodo del trie, eliminando la conversione di stringa per corrispondenza e la ricerca nell'hash-map laterale.

Ciclo interno:

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

Questo è lo stesso algoritmo, con la stessa forma del trie e lo stesso pattern di accesso, solo senza allocazioni e senza ricerche ridondanti. Risultati a 514 token:

Motore	p50	Istruzioni	IPC	Allocazioni
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
Baseline	156 µs	1,53M	2,62	0

Il baseline senza allocazioni confrontato con le implementazioni di tokenizzatore esistenti a 514 token.

La rimozione della materializzazione della stringa per corrispondenza e della ricerca secondaria nell'hash-map riduce le istruzioni di 2,4× rispetto al riferimento. L'IPC cala con esse (2,86 → 2,62) perché la caccia ai puntatori dell'hash-map, precedentemente nascosta dietro il lavoro dell'allocatore, è ora il costo per byte dominante.

Latenza p50 e p99 del tokenizzatore di riferimento e del baseline senza allocazioni in base alla dimensione dell'input.

Tasso di mancati colpi nella cache L2 del tokenizzatore di riferimento e del baseline senza allocazioni in base alla dimensione dell'input.

La baseline riduce la latenza p50 di 2-3× rispetto al riferimento a ogni dimensione e il suo tasso di miss L2 rimane vicino al 2-3% fino a 2k token contro l'8-48% del riferimento. Il costo rimanente è la caccia ai puntatori HashMap, che la sezione successiva affronta.

Il trie hash-map è una caccia ai puntatori. SP e IREE risolvono questo problema con un double-array trie (vedi sezione precedente) e prendiamo in prestito la stessa struttura dati. La differenza è che SP e IREE sono librerie per uso generale che supportano input in streaming, offset di output e modalità di codifica multiple; tale flessibilità costa loro una contabilità extra attorno al trie (un reticolo in SP, una pipeline a otto stadi in IREE). Non ne abbiamo bisogno, quindi inseriamo il trie direttamente nel ciclo di Viterbi e lasciamo perdere il resto.

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 passaggio di byte costa due letture di array, una somma di interi e un confronto. L'intero trie è di circa 9 MB, ma ogni singola codifica ne percorre solo ~100 KB, quindi i dati caldi rimangono in L2.

Motore	p50	p99	Istruzioni	L2 miss%	L3 miss%
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%
Baseline	155 µs	1.170 µs	1,53M	2,59%	16,66%
+ Darts	68 µs	80 µs	1,10M	6,45%	11,42%

Latenza del tokenizzatore, istruzioni e tassi di mancati colpi in cache dopo l'adozione del trie a doppio array.

Il wall-clock scende di 2,3× rispetto alla baseline e di 4,8× rispetto al riferimento. Le istruzioni scendono di un altro 28% rispetto alla baseline. Tagliare il reticolo e la pipeline ci permette di ritirare 1,10M di istruzioni per encode contro 1,82M di SP e 1,61M di IREE con la stessa forma del trie; il divario è la contabilità che mantengono per le API di uso generale di cui non abbiamo bisogno.

Il trie Darts dell'Ottimizzazione 1 necessita di due carichi dipendenti per ogni passaggio di byte. Prima l'offset base del genitore, poi l'entrata check nello slot figlio candidato per confermare che il genitore lo possiede. L'entrata check valida solo la proprietà. La stessa domanda può trovare risposta in una bitmap per nodo che registra quali byte hanno figli. Per la costruzione del double-array trie, il bit per il byte b nella bitmap del genitore è impostato esattamente quando lo slot candidato per quel byte è posseduto dal genitore. La bitmap e l'array check portano la stessa informazione, quindi possiamo eliminare completamente l'array check dalla struttura runtime.

La ricerca della bitmap si compila in un singolo test di bit contro una parola a 64 bit, che viene eseguito in circa un ciclo. La relativa istruzione POPCNT, che conta tutti i bit impostati in un registro, è ciò che altri design di trie indicizzati da bitmap usano per calcolare direttamente un offset figlio dalla bitmap. Non ne abbiamo bisogno perché la costruzione Darts memorizza già quell'offset nel campo base del nodo.

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

Mentre ci siamo, raggruppiamo i quattro campi per nodo (bitmap, base, ID token e punteggio) in una cache line da 64 byte. Una singola visita al nodo ora trascina dentro tutto: il punteggio e l'ID del token che ci servono nei nodi terminali, più la bitmap e la base per il passaggio successivo. L'array check viene usato solo durante la costruzione e viene eliminato dal layout runtime. L'indicizzazione del percorso critico non è controllata perché la costruzione del trie e la struttura del ciclo garantiscono già i limiti. Il compromesso è lo Spazio: a 64 byte per nodo × ~780K slot, il trie cresce da ~9 MB a ~50 MB.

Risultati a 514 token:

Motore	p50	p99	Istruzioni	Accessi L2	L2 miss%	L3 miss%
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%
Baseline	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%
+ Bitmap e Inline Packing	65 µs	77 µs	1,06M	1.8K	20,53%	15,01%

Risultati del tokenizzatore a 514 token dopo l'aggiunta del bitmap e del packing inline.

Il wall-clock scende sistematicamente di un altro 4,5% a 514 token e cresce con la lunghezza dell'input, raggiungendo il 6-7% a 4k e 16k. Eliminare l'array check trasforma il passaggio del trie per byte da due caricamenti di cache-line in uno, che è 2,5× meno accessi L2 per encode (1.8K contro i 4.6K del Darts semplice). Il tasso di miss salta dal 6% al 20%, ma la maggior parte è dovuta al denominatore: con meno della metà degli accessi, i miss che abbiamo occupano una quota maggiore.

I dati del trie sono di circa 50 MB (780.000 nodi a 64 byte ciascuno). In un processo Linux predefinito, quella memoria è supportata da pagine da 4 KB, quindi l'intero trie copre circa 12.000 pagine virtuali distinte, ben oltre ciò che il Translation Lookaside Buffer di primo livello della CPU può contenere (Sapphire Rapids ha 96 entrate). Ogni passaggio di Viterbi tocca un nodo diverso e su una codifica di 512 token bruciamo circa 9.000 cicli in camminate della page-table, circa il 3% del budget per codifica.

La soluzione è supportare il trie con huge pages da 2 MB invece di quelle da 4 KB. Gli stessi 50 MB ora coprono 25 pagine, si adattano interamente al Translation Lookaside Buffer e le directory della page-table stesse smettono di competere con il trie per la cache L2.

Linux ha due modi per supportare la memoria con pagine da 2 MB. Quello opportunistico è le transparent huge pages (THP): il tuo programma chiede al kernel di aggiornare una regione di memoria da 4 KB a pagine da 2 MB e un thread del kernel in background cerca di farlo quando ha tempo e memoria fisica contigua disponibile. L'altra alternativa è allocare il trie direttamente da un pool di huge-page pre-riservato che il kernel mette da parte all'avvio:

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

Il pool deve esistere sull'host (vm.nr_hugepages impostato all'avvio). Sulle nostre istanze di produzione sono riservate 10.561 huge pages; il trie ne consuma 24. Dopo la mmap, /proc/<pid>/smaps conferma la regione: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

Il guadagno è del 3–12% di wall-clock attraverso le dimensioni dell'input, coerente e proporzionale alla lunghezza del set di lavoro:

Token	Senza huge pages	Con huge pages	Δ
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%

Latenza del tokenizzatore con e senza huge pages su diverse dimensioni di input.

La vittoria più grande è a 4k token, dove il set di lavoro è abbastanza grande che il traffico della page-table era in competizione attiva con i dati del trie per la larghezza di banda L2. Oltre i 4k il guadagno si riduce perché gli L3 miss dominano e le camminate della page-table diventano una quota minore del tempo totale.

Mettendo insieme le tre ottimizzazioni, ecco il motore finale misurato rispetto a ogni altro tokenizer RoBERTa che abbiamo testato. Accanto ai binari nativi C++ SentencePiece e C IREE, abbiamo anche effettuato il benchmark dei crate wrapper Rust off-the-shelf a cui un utente di produzione si rivolgerebbe, sentencepiece e iree-tokenizer. Il crate sentencepiece misura una latenza superiore di 1,6-1,9× rispetto al binario C++ nativo a ogni dimensione dell'input, con il sovraccarico derivante dal wrapper che alloca una nuova lista di pezzi di token, ciascuno che detiene la propria copia del testo del pezzo, a ogni codifica. Il crate iree-tokenizer dimensiona il suo buffer interno a ogni input, il che evita una costante allocazione per-codifica che il binario IREE nativo paga indipendentemente dalla dimensione dell'input, e questo conta di più agli input brevi dove l'allocazione altrimenti dominerebbe il lavoro di tokenizzazione effettivo. Alle lunghezze dell'input di produzione il costo costante si ammortizza e il sovraccarico per chiamata del wrapper derivante dalla copia dei token di nuovo in Rust spinge la sua latenza sopra quella del binario nativo.

Latenza p50 e p99 per encode tra le implementazioni del tokenizzatore e le dimensioni di input.

Il motore finale esegue lo stesso Viterbi sullo stesso vocabolario da 250.000 token, non alloca nulla sul percorso critico e produce un output identico ai token rispetto al riferimento. Ogni passaggio di byte è ora un singolo caricamento di cache-line da 64 byte, un test di bitmap e una somma di interi, in calo rispetto a una sonda HashMap, una chiamata all'allocatore e diversi branch per casi speciali nell'implementazione della baseline. Le istruzioni per codifica sono scese da 3,66M a 1,04M, un calo di 3,5×. Il wall-clock corrisponde a quel rapporto agli input brevi e si allarga ulteriormente agli input lunghi, dove le allocazioni per token del riferimento iniziano a superare L2 e L3.

In pratica questo lavoro ha ridotto l'utilizzo della CPU nello strato di pre-elaborazione del nostro stack di inferenza di 5-6× e ha ridotto la latenza del reranker di decine di millisecondi.

La maggior parte della riduzione della latenza è derivata da modifiche leggere che hanno rimosso lavoro sprecato: abbandonare le allocazioni per codifica, sostituire il trie HashMap-children con un double-array e rimuovere il sovraccarico dell'API generica. L'uso della bitmap per nodo esistente come controllo di validità e il supporto del trie con huge pages da 2 MB hanno contribuito insieme per il 10-20% in aggiunta al passaggio più ampio a un double-array trie, a seconda della dimensione dell'input.

I crate wrapper Rust off-the-shelf attorno a SentencePiece e IREE illustrano lo stesso punto al contrario. Sopra i tokenizer nativi C e C++, il confine FFI e le allocazioni per chiamata che i wrapper introducono sono visibili nei numeri della latenza. Il sorgente Rust per l'encoder Viterbi finale è disponibile su 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}
}

Motivazione

L'algoritmo di Viterbi

Implementazione di riferimento

Una nota sulla pre-elaborazione

Dove la implementazione di riferimento spende il suo tempo

Come abbiamo misurato

Perché SentencePiece e IREE sono più veloci: il Double-Array Trie

Baseline: Un porting a zero allocazioni del riferimento

Ottimizzazione 1: Double-Array Trie senza astrazioni extra

Ottimizzazione 2: Bitmap e compressione inline

Ottimizzazione 3: Pagine enormi per il trie

Confronto finale

Pensieri conclusivi

Riferimenti

Citazione