Mellora do rendemento da CPU do tokenizador Unigram

Volvemos a implementar o noso tokenizador Unigram desde cero como un proxecto centrado no rendemento. Con lonxitudes de entrada de produción, o novo codificador reduce a latencia p50 nun 5× aproximadamente en comparación coa caixa de tokenizadores de Hugging Face, ~2× en comparación con SentencePiece (C++) e ~1,5× en comparación co tokenizador de IREE (C), sen ningunha asignación de montón en estado estacionario. En produción, isto reduciu a utilización da CPU na nosa pila de inferencia entre 5 e 6 veces e recortou milisegundos de dous díxitos na latencia do re-clasificador.

O custo de inferencia de LLM adoita formularse como traballo de GPU: cachés KV, núcleos de atención, encamiñamento de expertos. Mais unha parte significativa do tráfico de produción execútase en modelos moito máis pequenos, de dous a tres ordes de magnitude inferiores aos transformers de vangarda. Os modelos de incrustación, os clasificadores e os re-clasificadores seguen executándose en GPUs, mais a pasada de avance é o suficientemente curta como para que todo o que a rodea comece a importar.

O caso máis claro é un re-clasificador que puntúa centos de documentos candidatos por solicitude. Cun modelo pequeno, o cálculo da GPU adoita rematar en milisegundos dun só díxito, mais cada entrada aínda ten que pasar primeiro pola tokenización do lado da CPU, convertendo o texto en IDs de vocabulario. Cando o lote é grande, iso pode representar unha fracción significativa da latencia total da solicitude.

Este traballo céntrase en XLM-RoBERTa, un vocabulario Unigram de 250 mil tokens adestrado con SentencePiece. Os codificadores da familia RoBERTa axustados con finura son unha opción común para tarefas de clasificación, recuperación e simitude. Informamos da latencia por codificación dun só fío en todo momento; en produción paralelizamolos entre núcleos usando rayon.

Ao longo desta publicación, tratamos a caixa de tokenizadores de Hugging Face como a referencia coa que realizamos as nosas probas comparativas, xa que é o tokenizador predeterminado en Rust ao que acoden a maioría dos equipos. SentencePiece (C++) e IREE (C) inclúense para comparación como as outras dúas principais implementacións de produción.

A tokenización Unigram, introducida por Kudo en 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) e incluída en SentencePiece, formula a segmentación como un problema de ruta máis probable: cada token do vocabulario ten unha log-probabilidade aprendida, e o codificador escolle a segmentación cuxas puntuacións de token suman o valor máis alto. O problema resólvese co algoritmo de Viterbi (1967), empregando as posicións de bytes como capas do grafo e os tokens de vocabulario como arestas que abranguen un rango de bytes contiguos. A recorrencia é unha programación dinámica (DP) unidimensional sobre as posicións 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

O bucle exterior é lineal respecto á lonxitude da entrada, limitado en cada posición polo token de vocabulario máis longo (48 bytes UTF-8 para XLM-RoBERTa). O percorrido interior do trie é a ruta crítica: nunha entrada de 16 mil tokens, executa centos de miles de transicións e retira decenas de millóns de instrucións por codificación.

Os tokenizadores de Hugging Face v0.22 implementan a recorrencia directamente, cunha táboa hash para a busca inversa e fillos de táboa hash en cada nodo do 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 o bucle interior (simplificado) — cada coincidencia de busca de prefixo asigna unha nova cadea e despois proba a táboa 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
}

O código está en Rust lexible, mais realízase moito traballo que é invisible a nivel de código fonte. A seguinte sección desglosa a onde vai o tempo.

Antes de que se execute Viterbi, a entrada pasa pola normalización Unicode (NFKC) e pola marcación de límites de palabras. A profilaxe sitúa o preprocesamento en menos do 5% do tempo total de codificación nos nosos escenarios, polo que nos centramos na pasada de avance de Viterbi para o resto da publicación.

Hai dúas formas de acelerar o código de CPU: executar menos instrucións ou acadar un IPC superior. Cachegrind mide a primeira, perf stat mide a segunda. Cada optimización desta publicación recae nunha ou noutra.

As CPUs sofren bloqueos por catro motivos: fallos de caché (L1d ~4 ciclos, un fallo de L3 a DRAM son centos), cadeas de carga dependentes que o motor fóra de orde non pode paralelizar, erros de predición de saltos (~15-20 ciclos descartados por cada suposición incorrecta) e chamadas ao asignador. A implementación do motor unigram de referencia atópase cos catro:

Gargalo	Mecanismo	Impacto medido
Asignación por coincidencia	String::from_utf8 nos bytes coincidentes máis unha busca en AHashMap por cada coincidencia no trie; o ID do token xa está na folla, mais a referencia non o aloxa aí	7.295 asignacións con 512 tokens, 14.542 con 1k, 299.171 con 16k
Busca de punteiros por byte	Un AHashMap en cada nodo do trie, aplicando hash ao seguinte byte e desreferenciando dous punteiros por paso; a estrutura Node tamén agrupa campos fríos como is_leaf, arrastrándoos á caché en cada visita	A taxa de fallos en L1d mantense baixa (~2,5% a 1k) porque as páxinas quentes do asignador viven en L1; o custo maniféstase como latencia de carga dependente
Saturación de L2 en entradas longas	A táboa de DP, o borrador de prefixos e a lista de tokens de saída asígnanse de novo en cada chamada	A taxa de fallos en L2 sube do 12% con 128 tokens ao 50% con 16k a medida que a pegada de asignación acumulada supera a caché L2 de 2 MB por núcleo

Referencia dos pescozos de botella do tokenizador, mecanismos e impactos medidos.

O custo dominante é a asignación por coincidencia. Con 1k tokens, a implementación de referencia retira (completa) aproximadamente 7M de instrucións por codificación a un IPC de 2,8, a maioría delas en chamadas ao asignador e conversións de cadeas. Os fallos de L2 só importan a partir de ~1k tokens, unha vez que a pegada de asignación acumulada desborda a caché por núcleo.

Latencia p50 e p99 do tokenizador de referencia por tamaño de entrada.

Taxa de fallos da caché L2 do tokenizador de referencia por tamaño de entrada.

Bytes asignados por codificación polo tokenizador de referencia a través de distintos tamaños de entrada.

A asignación por token é constante: uns 2 KB e ~18 asignacións por token, independentemente do tamaño da entrada. A non linealidade da latencia provén de L2: os bytes acumulados superan os 2 MB arredor de 1k tokens, a taxa de fallos sube do 8% a ~50%, e a curva p99 flexiónase para coincidir.

Comparamos tres tokenizadores de produción, todos executando o mesmo algoritmo Unigram co mesmo modelo XLM-RoBERTa:

Motor	Lingua	Trie
Hugging Face	Rust	Trie de prefixos con fillos HashMap
SentencePiece	C++	Trie de dobre matriz Darts
Tokenizador IREE	C	Trie de dobre matriz Darts

Implementacións do tokenizador, linguaxes e representacións trie usadas na comparación.

Vinculamos SentencePiece e IREE mediante FFI directo de C/C++ en lugar dos seus encapsuladores en Rust; os encapsuladores engaden asignacións por chamada e conversións de cadeas que non forman parte do custo da biblioteca subxacente.

Todos os motores constrúense coas optimizacións completas (O3, orientación nativa, optimización en tempo de enlace). As medicións executanse fixadas a un núcleo silencioso nun Intel Xeon Platinum 8488C. Para cada motor e tamaño de entrada recollemos a latencia de reloxo de parede (p50 e p99, 10.000 iteracións despois de 1.000 de quecemento), instrucións retiradas mediante cachegrind de Valgrind, taxas de IPC e fallos de caché dos contadores de hardware de perf stat, e asignacións de montón contadas por un encapsulador de asignador global arredor da xanela medida.

Con 514 tokens (os +2 sobre 512 son a inxección de BOS/EOS):

Implementación	Latencia p50	Instrucións	IPC	Asignacións
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

Latencia do tokenizador, número de instrucións, IPC e asignacións en 514 tokens.

SP e IREE obteñen aproximadamente un terzo da latencia da referencia. Retiran preto da metade de instrucións, executánse a un IPC superior e asignan dunha a dúas ordes de magnitude menos.

SentencePiece e IREE empregan o trie de dobre matriz introducido por Aoe en 1989 (Aoe, 1989; Karoonboonyanan). O trie codifícase en dúas matrices de enteiros planas, 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

Unha busca de fillo é:

cpp

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

Cada paso do percorrido do trie só precisa de dúas lecturas de matriz, unha suma de enteiros e unha comparación, sen hashing, sen busca de punteiros e sen asignación de montón por nodo. A construción coloca cada nodo de modo que os seus fillos caian en ranuras libres na matriz de validación, permitindo que varios pais compartan rexións sempre que os seus fillos difiran no byte que seguen.

Para o vocabulario de 250K de XLM-RoBERTa, todo o trie encadea en ~9 MB de memoria contigua (~780K ranuras × 12 bytes para o desprazamento, a validación e o ID de token por ranura). O conxunto de traballo en quente dunha soa codificación é da orde de 100 KB, cómodamente en L2, e o prefetcher do hardware bloquéase no patrón de acceso regular.

Pola contra, a referencia mantén unha táboa hash de fillos por nodo, o que custa aproximadamente catro cargas dependentes por paso de byte. Aplícase hash ao byte, cárgase o cubo, séguese un punteiro á entrada do cubo, séguese outro punteiro ao nodo fillo no montón e despois consúltase a súa propia táboa hash para o seguinte byte. Os obxectos do montón dispérsanse e a maioría das cargas fallan na caché.

Tanto SP como IREE almacenan o ID do token directamente no nodo do trie. A implementación do trie de referencia só devolve os bytes coincidentes, polo que o codificador ten que asignar unha nova cadea para cada coincidencia, buscala nunha táboa hash lateral para recuperar o ID do token e ler a puntuación dun vector separado. SP devolve o ID como o valor almacenado do trie e le a puntuación dunha matriz plana de flotantes:

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 emprega unha API de cursor con estado en lugar de devolver todas as coincidencias nunha soa chamada, mais o modelo de custos é o mesmo: os IDs de tokens e as puntuacións son accesibles como índices de matriz plana, sen asignación nin hashing no propio percorrido do trie. IREE asigna só un búfer de saída por codificación; ademais, SP asigna un nodo de retícula por coincidencia durante a DP, sumando aproximadamente 1.559 asignacións con 514 tokens.

Antes de cambiar o trie, illamos o custo do algoritmo do custo de implementación: a que velocidade se executa a referencia co traballo innecesario eliminado?

A liña base mantén o trie de fillos con táboa hash da referencia mais fai dous cambios. Unha estrutura de borrador propiedade da persoa chamada contén a táboa de DP, os bytes preprocesados e os IDs de tokens de saída, polo que a ruta crítica non realiza asignacións despois do quecemento. O ID e a puntuación do token almacénanse directamente no nodo do trie, eliminando a conversión de cadeas por coincidencia e a busca na táboa 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,
        }
    }
}

Este é o mesmo algoritmo, coa mesma forma de trie e o mesmo patrón de acceso, só que sen asignacións e sen buscas redundantes. Resultados con 514 tokens:

Motor	p50	Instrucións	IPC	Asignacións
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
Liña base	156 µs	1,53M	2,62	0

A liña base sen asignacións comparada coas implementacións de tokenizador existentes en 514 tokens.

A eliminación da materialización de cadeas por coincidencia e da busca secundaria na táboa hash reduce as instrucións 2,4× en comparación coa referencia. O IPC cae con elas (2,86 → 2,62) porque a busca de punteiros na táboa hash, antes oculta tras o traballo do asignador, é agora o custo dominante por byte.

Latencia p50 e p99 do tokenizador de referencia e da liña base sen asignacións por tamaño de entrada.

Taxa de fallos da caché L2 do tokenizador de referencia e da liña base sen asignacións por tamaño de entrada.

A liña base reduce a latencia p50 nun factor de 2-3× en comparación coa referencia en calquera tamaño, e a súa taxa de fallos en L2 manténse preto do 2-3% ata 2k tokens fronte ao 8-48% da referencia. O custo restante é a busca de punteiros en HashMap, que a seguinte sección aborda.

O trie de táboa hash é unha busca de punteiros. SP e IREE resolvem isto cun trie de dobre matriz (véxase a sección anterior) e tomamos prestada a mesma estrutura de datos. A diferenza é que SP e IREE son bibliotecas de propósito xeral que admiten entradas en fluxo, desprazamentos de saída e múltiples modos de codificación; esa flexibilidade cóstasles un rexistro adicional arredor do trie (unha retícula en SP, unha canalización de oito etapas en IREE). Non necesitamos nada diso, polo que incrustamos o trie directamente no bucle de Viterbi e prescindimos do 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 paso de byte custa dúas lecturas de matriz, unha suma de enteiros e unha comparación. Todo o trie ten uns 9 MB, mais calquera codificación individual só percorre ~100 KB del, polo que os datos quentes permanecen en L2.

Motor	p50	p99	Instrucións	% fallos L2	% fallos 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%
Liña base	155 µs	1.170 µs	1,53M	2,59%	16,66%
+ Darts	68 µs	80 µs	1,10M	6,45%	11,42%

Latencia do tokenizador, instrucións e taxas de fallos de caché despois de adoptar o trie de matriz dobre.

O reloxo de parede cae 2,3× desde a liña base e 4,8× desde a referencia. As instrucións baixan outro 28% desde a liña base. A eliminación da retícula e da canalización permítenos retirar 1,10M de instrucións por codificación fronte aos 1,82M de SP e os 1,61M de IREE coa mesma forma de trie; a diferenza é o rexistro que manteñen para as APIs de propósito xeral que non necesitamos.

O trie Darts da Optimización 1 precisa dúas cargas dependentes por paso de byte. Primeiro o desprazamento base do pai, despois a entrada check na ranura do fillo candidato para confirmar que o pai é o propietario. A entrada check só valida a propiedade. A mesma pregunta pode ser respondida cun mapa de bits por nodo que rexistra qué bytes teñen fillos. Mediante a construción do trie de dobre matriz, o bit para o byte b no mapa de bits do pai establécese exactamente cando a ranura candidata para ese byte é propiedade do pai. O mapa de bits e a matriz check levan a mesma información, polo que podemos eliminar por completo a matriz check da estrutura en tempo de execución.

A busca no mapa de bits compílase nunha soa proba de bit contra unha palabra de 64 bits, o que se executa aproximadamente nun ciclo. A instrución POPCNT relacionada, que conta todos os bits establecidos nun rexistro, é a que outros deseños de trie indexados por mapa de bits usan para calcular un desprazamento de fillo directamente desde o mapa de bits. Non o necesitamos porque a construción Darts xa almacena ese desprazamento no campo base do 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)
}

Xa que estamos, empaquetamos os catro campos por nodo (mapa de bits, base, ID de token e puntuación) nunha única liña de caché de 64 bytes. Unha visita a un único nodo trae agora todo: a puntuación e o ID do token que necesitamos nos nodos terminais, ademais do mapa de bits e a base para o seguinte paso. A matriz check só se usa durante a construción e elimínase da disposición en tempo de execución. A indexación da ruta crítica non está verificada porque a construción do trie e a estrutura do bucle xa garanten os límites. O compromiso é o espazo: a 64 bytes por nodo × ~780K ranuras, o trie crece de ~9 MB a ~50 MB.

Resultados con 514 tokens:

Motor	p50	p99	Instrucións	Accesos L2	% fallos L2	% fallos 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%
Liña 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 e empaquetado en liña	65 µs	77 µs	1,06M	1,8K	20,53%	15,01%

Resultados do tokenizador en 514 tokens despois de engadir bitmap e empaquetado en liña.

O reloxo de parede cae sistematicamente outro 4,5% con 514 tokens e medra coa lonxitude da entrada, acadando o 6-7% con 4k e 16k. A eliminación da matriz check converte o paso do trie por byte de dúas cargas de liña de caché nunha soa, o que supón 2,5× menos accesos a L2 por codificación (1,8K fronte aos 4,6K do Darts simple). A taxa de fallos salta do 6% ao 20%, mais gran parte diso débese ao denominador: con menos da metade de accesos, os fallos que temos ocupan unha maior proporción.

Os datos do trie ocupan uns 50 MB (780.000 nós a 64 bytes cada un). Nun proceso Linux predeterminado, esa memoria está respaldada por páxinas de 4 KB, polo que o trie completo abrangue preto de 12.000 páxinas virtuais distintas, moi por riba do que o búfer de tradución lookaside (TLB) de datos de primeiro nivel da CPU pode conter (Sapphire Rapids ten 96 entradas). Cada paso de Viterbi toca un nodo diferente e, nunha codificación de 512 tokens, consumimos aproximadamente 9.000 ciclos en percorridos pola táboa de páxinas, arredor do 3% do orzamento por codificación.

A solución consiste en respaldar o trie con páxinas xigantes (huge pages) de 2 MB en lugar das de 4 KB. Os mesmos 50 MB abarcan agora 25 páxinas, encaixan por completo no TLB, e os propios directorios da táboa de páxinas deixan de competir co trie pola caché L2.

Linux ten dúas formas de respaldar a memoria con páxinas de 2 MB. A oportunista son as páxinas xigantes transparentes (THP): o programa pide ao núcleo que anovar unha rexión de memoria de páxinas de 4 KB a 2 MB, e un fío en segundo plano do núcleo tenta facelo cando ten tempo e memoria física contigua dispoñible. A outra alternativa é asignar o trie directamente desde un grupo de páxinas xigantes pre-reservadas que o núcleo aparta no arranque:

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

O grupo debe existir no host (vm.nr_hugepages configurado no arranque). Nas nosas instancias de produción, resérvanse 10.561 páxinas xigantes; o trie consume 24 delas. Despois do mmap, /proc/<pid>/smaps confirma a rexión: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

A ganancia é dun 3-12% no reloxo de parede entre tamaños de entrada, de forma consistente e proporcional á lonxitude do conxunto de traballo:

Tokens	Sen páxinas xigantes	Con páxinas xigantes	Δ
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%

Latencia do tokenizador con e sen páxinas enormes a través de distintos tamaños de entrada.

A maior vantaxe obtense con 4k tokens, onde o conxunto de traballo é o suficientemente grande como para que o tráfico da táboa de páxinas estivese a competir activamente cos datos do trie polo ancho de banda de L2. Máis alá de 4k a ganancia redúcese porque os fallos de L3 dominan e os percorridos pola táboa de páxinas pasan a representar unha menor proporción do tempo total.

Xuntando as tres optimizacións, aquí está o motor final medido fronte a calquera outro tokenizador de RoBERTa que probamos. Xunto cos binarios nativos en C++ de SentencePiece e en C de IREE, tamén avaliamos as caixas de encapsulamento en Rust dispoñibles que un usuario de produción empregaría, sentencepiece e iree-tokenizer. A caixa sentencepiece rexistra unha latencia de 1,6-1,9× superior á do binario nativo en C++ en cada tamaño de entrada, cunha sobrecarga derivando de que o encapsulador asigna unha nova lista de pezas de tokens, cada unha mantendo a súa propia copia do texto da peza, en cada codificación. A caixa iree-tokenizer axusta o seu búfer interno a cada entrada, o que evita unha asignación constante por codificación que o binario nativo de IREE paga independentemente do tamaño da entrada, e iso importa máis en entradas curtas onde a asignación doutro xeito dominaría o traballo real de tokenización. Con lonxitudes de entrada de produción o custo constante amortízase e a propia sobrecarga por chamada do encapsulador ao copiar os tokens de volta a Rust eleva a súa latencia por riba da do binario nativo.

Latencia p50 e p99 por codificación entre implementacións de tokenizador e tamaños de entrada.

O motor final executa o mesmo Viterbi sobre o mesmo vocabulario de 250.000 tokens, non asigna nada na ruta crítica e produce unha saída exacta en canto a tokens en comparación coa referencia. Cada paso de byte é agora unha única carga de liña de caché de 64 bytes, unha proba de mapa de bits e unha suma de enteiros, fronte a unha proba de HashMap, unha chamada ao asignador e varias ramificacións de casos especiais na implementación de base. As instrucións por codificación baixaron de 3,66M a 1,04M, unha caída de 3,5×. O reloxo de parede coincide con esa proporción en entradas curtas e amplíase aínda máis en entradas longas, onde as asignacións por token da referencia comezan a desbordar L2 e L3.

Na práctica, este traballo reduciu a utilización da CPU na capa de preprocesamento da nosa pila de inferencia entre 5 e 6 veces e recortou milisegundos de dous díxitos na latencia do re-clasificador.

A maior parte da redución de latencia provén de cambios lixeiros que eliminaron o traballo desmerecido: eliminar as asignacións por codificación, substituír o trie de fillos con HashMap por un de dobre matriz e eliminar a sobrecarga da API xenérica. O uso do mapa de bits por nodo existente como comprobación de validez e o respaldo do trie con páxinas xigantes de 2 MB combinados supuxo un 10-20% ademais do cambio maior a un trie de dobre matriz, dependendo do tamaño da entrada.

As caixas de encapsulamento en Rust dispoñibles ao redor de SentencePiece e IREE ilustran o mesmo punto á inversa. Ademais dos tokenizadores nativos en C e C++, a fronteira FFI e as asignacións por chamada que introducen os encapsuladores son visibles nos números de latencia. O código fonte en Rust para o codificador de Viterbi final está dispoñible en github.com/perplexityai/pplx-garden.

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

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

Tokenizador 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áxinas xigantes — 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ón

O algoritmo de Viterbi

Implementación de referencia

Unha nota sobre o preprocesamento

Onde gasta o seu tempo a implementación de referencia

Como medimos

Por que SentencePiece e IREE son máis rápidos: o trie de dobre matriz

Liña base: unha adaptación sen asignacións da referencia

Optimización 1: trie de dobre matriz sen abstraccións adicionais

Optimización 2: Mapa de bits e empaquetado en liña

Optimización 3: Páxinas grandes para o trie

Comparación final

Reflexións finais

Referencias

Cita