Perplexity AI Lancia il Tokenizer Rust Open Source: 5 Volte Più Veloce di Hugging Face

Il team di ricerca di Perplexity AI ha annunciato la pubblicazione in open source di una nuova implementazione del tokenizer Unigram scritta interamente in Rust. Questo componente, parte del repository di tecnologie di inferenza pplx-garden, promette di rivoluzionare le prestazioni della CPU nei carichi di lavoro legati ai modelli linguistici (LLM), riducendo la latenza p50 di circa 5 volte rispetto allo standard di Hugging Face e abbattendo drasticamente il consumo di risorse computazionali.

Il collo di bottiglia invisibile dell'inferenza

Sebbene l'attenzione del settore sia quasi sempre focalizzata sull'ottimizzazione dei kernel GPU e delle memorie KV cache, i modelli più piccoli — come quelli utilizzati per embedding, classificazione e reranking — presentano una dinamica differente. Questi modelli sono spesso da due a tre ordini di grandezza più piccoli rispetto ai transformer di frontiera.

In uno scenario in cui un reranker deve valutare centinaia di documenti candidati per ogni singola richiesta, il calcolo sulla GPU può concludersi in pochi millisecondi. Tuttavia, ogni input deve prima passare attraverso la tokenizzazione lato CPU. Quando le dimensioni dei batch sono elevate, la tokenizzazione smette di essere un'operazione trascurabile e diventa una frazione significativa della latenza totale della richiesta.

Il lavoro di Perplexity si concentra in particolare su XLM-RoBERTa, un modello con un vocabolario Unigram da 250.000 token addestrato con SentencePiece, ampiamente utilizzato in produzione per attività di ranking e similarità.

Cos'è la tokenizzazione Unigram?

Introdotta da Kudo nel 2018, la tokenizzazione Unigram modella la segmentazione del testo come un problema del "percorso più probabile". Ad ogni token del vocabolario è assegnata una log-probabilità appresa; il tokenizer sceglie quindi la segmentazione in cui la somma dei punteggi dei token è massima.

Per trovare questo percorso ottimale viene impiegato l'algoritmo di Viterbi, una tecnica di programmazione dinamica del 1967. In questo contesto, le posizioni dei byte formano i livelli di un grafo, mentre i token del vocabolario rappresentano gli archi che coprono intervalli di byte contigui.

Analisi delle inefficienze: il confronto con Hugging Face

Il crate tokenizers di Hugging Face è lo standard di riferimento in ambiente Rust, ma l'analisi di Perplexity ha evidenziato tre pattern estremamente costosi nelle performance:

A una lunghezza di 514 token, l'implementazione di riferimento soffre di un sovraccarico costante: circa 2 KB e 18 allocazioni per ogni singolo token, indipendentemente dalla dimensione dell'input. Questo porta a una degradazione severa quando gli input lunghi superano la capacità della cache L2 del core CPU.

Le tre ottimizzazioni chiave di Perplexity

1. Double-Array Trie

Mentre Hugging Face memorizza i figli in una HashMap per ogni nodo, richiedendo calcoli hash e de-referenziazioni multiple, Perplexity ha implementato un double-array trie (struttura introdotta da Aoe nel 1989). Questa struttura codifica l'intero albero in due array di interi piatti (base e check). La ricerca di un figlio diventa una semplice operazione aritmetica:

rustCopia
// Pseudo-logica di ricerca
let next = base[node] + byte;
if check[next] == node {
    // Transizione valida
}

Il risultato è drastico: la latenza p50 è scesa da 155 µs (ottenuti con un porting a zero allocazioni) a soli 68 µs.

2. Bitmap e Inline Packing

Per eliminare anche il secondo caricamento dall'array (il controllo check), Perplexity ha introdotto una bitmap per nodo (32 byte) che indica quali dei 256 possibili byte hanno transizioni valide. Hanno inoltre impacchettato tutti i dati del nodo (bitmap, base, ID token e score) in una singola linea di cache da 64 byte. In questo modo, un singolo caricamento dalla memoria fornisce tutte le informazioni necessarie per lo step successivo.

3. Huge Pages (2 MB)

Con un trie di circa 50 MB, l'utilizzo di pagine standard da 4 KB comporterebbe circa 12.000 pagine virtuali, saturando il TLB (Translation Lookaside Buffer) della CPU e causando continui "page-table walks". Utilizzando le Huge Pages da 2 MB tramite mmap, l'intera struttura risiede in sole 25 pagine, eliminando quasi totalmente i ritardi di traduzione degli indirizzi.

Risultati dei Benchmark

I test sono stati eseguiti su un core singolo di un processore Intel Xeon Platinum 8488C. I dati mostrano una superiorità netta dell'implementazione di Perplexity rispetto a tutte le alternative esistenti.

L'efficienza si riflette anche nel numero di istruzioni completate: l'implementazione finale richiede solo 1,04 milioni di istruzioni per un input da 514 token, contro i 3,60 milioni di Hugging Face.