Mejorando el rendimiento de la CPU del tokenizador Unigram

Hemos reimplementado nuestro tokenizador Unigram desde cero como un proyecto enfocado en el rendimiento. Con longitudes de entrada de producción, el nuevo codificador reduce la latencia p50 aproximadamente 5 veces frente a la caja de tokenizadores de Hugging Face, ~2 veces frente a SentencePiece (C++) y ~1.5 veces frente al tokenizador de IREE (C), con cero asignaciones de montón en estado estable. En producción, esto redujo la utilización de CPU en nuestra pila de inferencia entre 5 y 6 veces y recortó milisegundos de dos dígitos de la latencia del reranker.

El costo de inferencia de LLM generalmente se formula como trabajo de GPU: cachés KV, núcleos de atención, enrutamiento de expertos. Pero una parte significativa del tráfico de producción se ejecuta en modelos mucho más pequeños, de dos a tres órdenes de magnitud más pequeños que los transformadores de vanguardia. Los modelos de incrustación, clasificadores y rerankers aún se ejecutan en GPUs, pero el paso hacia adelante es lo suficientemente corto como para que todo lo que lo rodea comience a importar.

El caso más claro es un reranker que califica cientos de documentos candidatos por solicitud. Con un modelo pequeño, el cómputo de GPU generalmente termina en milisegundos de un solo dígito, pero cada entrada aún tiene que pasar primero por la tokenización del lado de la CPU, convirtiendo el texto en IDs de vocabulario. Cuando el lote es grande, eso puede representar una fracción significativa de la latencia total de la solicitud.

Este trabajo tiene como objetivo XLM-RoBERTa, un vocabulario Unigram de 250K tokens entrenado con SentencePiece. Los codificadores de la familia RoBERTa ajustados son una opción común para tareas de clasificación, recuperación y similitud. Informamos la latencia por codificación de un solo hilo en todas partes; en producción paralelizamos a través de núcleos usando rayon.

A lo largo de esta publicación, tratamos la caja de tokenizadores de Hugging Face como la referencia con la que comparamos, ya que es el tokenizador Rust predeterminado al que recurren la mayoría de los equipos. SentencePiece (C++) e IREE (C) se incluyen para comparar como las otras dos implementaciones principales de producción.

La tokenización Unigram, introducida por Kudo en 2018 (Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates) e incluida en SentencePiece, enmarca la segmentación como un problema de ruta de mayor probabilidad: cada token de vocabulario tiene una log-probabilidad aprendida, y el codificador elige la segmentación cuya suma de puntajes de tokens es el valor más alto. El problema se resuelve con el algoritmo de Viterbi (1967), con posiciones de bytes como capas de grafo y tokens de vocabulario como bordes que abarcan un rango de bytes contiguo. La recurrencia es una DP unidimensional sobre las posiciones 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 externo es lineal en la longitud de entrada, limitado en cada posición por el token de vocabulario más largo (48 bytes UTF-8 para XLM-RoBERTa). El recorrido interno del trie es la ruta crítica: en una entrada de 16k tokens, ejecuta cientos de miles de transiciones y retira decenas de millones de instrucciones por codificación.

La versión 0.22 de tokenizadores de Hugging Face implementa la recurrencia directamente, con un mapa hash para la búsqueda inversa y niños-mapa-hash en cada nodo 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
}

Y el bucle interno (simplificado): cada coincidencia de búsqueda de prefijo asigna una cadena nueva y luego sondea el mapa 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 código es Rust legible, pero ocurre mucho trabajo que es invisible a nivel de fuente. La siguiente sección desglosa a dónde va el tiempo.

Antes de que se ejecute Viterbi, la entrada pasa por la normalización Unicode (NFKC) y el marcado de límites de palabras. El perfilado sitúa el preprocesamiento en menos del 5% del tiempo total de codificación en nuestros escenarios, por lo que nos centramos en el paso hacia adelante de Viterbi durante el resto de la publicación.

Hay dos formas de acelerar el código de CPU: ejecutar menos instrucciones u obtener un IPC más alto. Cachegrind mide lo primero, perf stat mide lo segundo. Cada optimización en esta publicación se aplica a una o a otra.

Las CPUs se bloquean por cuatro cosas: fallos de caché (L1d ~4 ciclos, un fallo L3-a-DRAM es de cientos), cadenas de carga dependientes que el motor fuera de orden no puede paralelizar, predicciones incorrectas de ramas (~15-20 ciclos descartados por suposición incorrecta) y llamadas al asignador. La implementación del motor unigram de referencia alcanza las cuatro:

Cuello de botella	Mecanismo	Impacto medido
Asignación por coincidencia	String::from_utf8 en los bytes coincidentes más una búsqueda en AHashMap por coincidencia trie; el ID del token ya está en la hoja, pero la referencia no lo almacena allí	7,295 asignaciones en 512 tokens, 14,542 en 1k, 299,171 en 16k
Persecución de punteros por byte	Un AHashMap en cada nodo trie, realizando hash del siguiente byte y desreferenciando dos punteros por paso; la estructura Node también agrupa campos fríos como is_leaf, arrastrándolos a la caché en cada visita	La tasa de fallos de L1d se mantiene baja (~2.5% en 1k) porque las páginas activas del asignador viven en L1; el costo aparece como latencia de carga dependiente
Agitación de L2 en entradas largas	La tabla DP, el borrador de prefijo y la lista de tokens de salida se asignan recién en cada llamada	La tasa de fallos de L2 sube del 12% en 128 tokens al 50% en 16k a medida que la huella de asignación acumulativa excede los 2 MB de L2 por núcleo

Consulte los cuellos de botella, los mecanismos y los impactos medidos del tokenizer de referencia.

El costo dominante es la asignación por coincidencia. A 1k tokens, la implementación de referencia retira (completa) aproximadamente 7M de instrucciones por codificación a un IPC de 2.8, la mayoría de ellas en llamadas al asignador y conversiones de cadenas. Los fallos de L2 solo importan después de ~1k tokens, una vez que la huella de asignación acumulativa desborda la caché por núcleo.

Latencia p50 y p99 del tokenizer de referencia por tamaño de entrada.

Tasa de fallos de caché L2 del tokenizer de referencia por tamaño de entrada.

Bytes asignados por codificación del tokenizer de referencia en distintos tamaños de entrada.

La asignación por token es constante: alrededor de 2 KB y ~18 asignaciones por token, independientemente del tamaño de la entrada. La no linealidad de la latencia proviene de L2: los bytes acumulados cruzan los 2 MB alrededor de 1k tokens, la tasa de fallos sube del 8% a ~50%, y la curva p99 se curva para coincidir.

Comparamos tres tokenizadores de producción, todos ejecutando el mismo algoritmo Unigram contra el mismo modelo XLM-RoBERTa:

Motor	Idioma	Trie
Hugging Face	Rust	Trie de prefijo con niños-HashMap
SentencePiece	C++	Trie de doble matriz Darts
Tokenizador IREE	C	Trie de doble matriz Darts

Implementaciones de tokenizer, lenguajes y representaciones de trie utilizadas en la comparación.

Vinculamos SentencePiece e IREE a través de FFI directo C/C++ en lugar de sus envoltorios Rust; los envoltorios añaden asignaciones por llamada y conversiones de cadenas que no son parte del costo de la biblioteca subyacente.

Todos los motores están construidos con optimizaciones completas (O3, orientación nativa, optimización en tiempo de enlace). Las mediciones se ejecutan fijadas a un núcleo silencioso en un Intel Xeon Platinum 8488C. Para cada motor y tamaño de entrada recopilamos latencia de reloj de pared (p50 y p99, 10,000 iteraciones después de 1,000 de calentamiento), instrucciones retiradas a través de Valgrind cachegrind, IPC y tasas de fallos de caché de los contadores de hardware perf stat, y asignaciones de montón contadas por un envoltorio de asignador global alrededor de la ventana medida.

A 514 tokens (el +2 sobre 512 es inyección BOS/EOS):

Implementación	Latencia p50	Instrucciones	IPC	Asignaciones
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 del tokenizer, recuento de instrucciones, IPC y asignaciones con 514 tokens.

SP e IREE se sitúan aproximadamente en un tercio de la latencia de referencia. Retiran aproximadamente la mitad de las instrucciones, funcionan a un IPC más alto y asignan de uno a dos órdenes de magnitud menos.

SentencePiece e IREE utilizan el trie de doble matriz introducido por Aoe en 1989 (Aoe, 1989; Karoonboonyanan). El trie está codificado en dos matrices de enteros planos, base y verificación:

cpp

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

Una búsqueda de niño es:

cpp

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

Cada paso del recorrido del trie necesita solo dos cargas de matriz, una suma de enteros y una comparación, sin hash, sin persecución de punteros y sin asignación de montón por nodo. La construcción coloca cada nodo de modo que sus hijos caigan en espacios libres en la matriz de validación, permitiendo que varios padres compartan regiones siempre que sus hijos difieran en el byte que siguen.

Para el vocabulario de 250K de XLM-RoBERTa, todo el trie cabe en ~9 MB de memoria contigua (~780K espacios × 12 bytes para el desplazamiento, validación e ID de token por espacio). El conjunto de trabajo activo de una sola codificación es del orden de 100 KB, cómodamente en L2, y el prefetcher de hardware se bloquea en el patrón de acceso regular.

La referencia, por el contrario, mantiene un mapa hash de hijos por nodo, lo que cuesta aproximadamente cuatro cargas dependientes por paso de byte. Hash del byte, cargar el cubo, seguir un puntero a la entrada del cubo, seguir otro puntero al nodo hijo en el montón, luego consultar su propio mapa hash para el siguiente byte. Los objetos del montón se dispersan y la mayoría de las cargas fallan en la caché.

SP e IREE almacenan el ID del token directamente en el nodo trie. La implementación del trie de referencia devuelve solo los bytes coincidentes, por lo que el codificador tiene que asignar una cadena nueva para cada coincidencia, buscarla en un mapa hash lateral para recuperar el ID del token y leer el puntaje de un vector separado. SP devuelve el ID como el valor almacenado del trie y lee el puntaje de una matriz de flotantes plana:

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 utiliza una API de cursor con estado en lugar de devolver todas las coincidencias en una sola llamada, pero el modelo de costos es el mismo: los IDs de los tokens y los puntajes son alcanzables como índices de matriz plana, sin asignación o hash en el recorrido del trie. IREE asigna solo un búfer de salida por codificación; SP asigna adicionalmente un nodo de red por coincidencia durante la DP, totalizando aproximadamente 1,559 asignaciones en 514 tokens.

Antes de cambiar el trie, aislamos el costo del algoritmo del costo de implementación: ¿qué tan rápido se ejecuta la referencia con el trabajo innecesario eliminado?

La línea base mantiene el trie de niños-mapa-hash de la referencia, pero hace dos cambios. Una estructura de borrador propiedad del llamador contiene la tabla DP, los bytes preprocesados y los IDs de los tokens de salida, por lo que la ruta crítica no realiza asignaciones después del calentamiento. El ID del token y el puntaje se almacenan directamente en el nodo trie, eliminando la conversión de cadena por coincidencia y la búsqueda en el mapa hash lateral.

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

Este es el mismo algoritmo, con la misma forma de trie y el mismo patrón de acceso, solo que sin asignaciones y sin búsquedas redundantes. Resultados en 514 tokens:

Motor	p50	Instrucciones	IPC	Asignaciones
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ínea base	156 µs	1.53M	2.62	0

La línea base sin asignaciones frente a las implementaciones de tokenizer existentes con 514 tokens.

Eliminar la materialización de cadenas por coincidencia y la búsqueda secundaria en el mapa hash reduce las instrucciones 2.4 veces frente a la referencia. El IPC cae con ellas (2.86 → 2.62) porque la persecución de punteros del mapa hash, previamente oculta detrás del trabajo del asignador, es ahora el costo dominante por byte.

Latencia p50 y p99 del tokenizer de referencia y de la línea base sin asignaciones por tamaño de entrada.

Tasa de fallos de caché L2 del tokenizer de referencia y de la línea base sin asignaciones por tamaño de entrada.

La línea base reduce la latencia p50 2-3 veces frente a la referencia en cada tamaño, y su tasa de fallos de L2 se mantiene cerca del 2-3% hasta 2k tokens frente al 8-48% de la referencia. El costo restante es la persecución de punteros del HashMap, que la siguiente sección aborda.

El trie de mapa hash es una persecución de punteros. SP e IREE resuelven esto con un trie de doble matriz (ver la sección anterior), y tomamos prestada la misma estructura de datos. La diferencia es que SP e IREE son bibliotecas de uso general que admiten entrada de streaming, desplazamientos de salida y múltiples modos de codificación; esa flexibilidad les cuesta contabilidad adicional alrededor del trie (una red en SP, un conducto de ocho etapas en IREE). No necesitamos nada de eso, así que colocamos el trie directamente en el bucle de Viterbi y descartamos el 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 cuesta dos lecturas de matriz, una suma de enteros y una comparación. Todo el trie es de alrededor de 9 MB, pero cualquier codificación individual solo recorre ~100 KB de él, por lo que los datos activos se mantienen en L2.

Motor	p50	p99	Instrucciones	% 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%
Línea base	155 µs	1,170 µs	1.53M	2.59%	16.66%
+ Darts	68 µs	80 µs	1.10M	6.45%	11.42%

Latencia del tokenizer, instrucciones y tasas de fallos de caché tras adoptar el trie de doble matriz.

El reloj de pared cae 2.3 veces desde la línea base y 4.8 veces desde la referencia. Las instrucciones caen otro 28% desde la línea base. Cortar la red y el conducto nos permite retirar 1.10M de instrucciones por codificación frente a los 1.82M de SP y los 1.61M de IREE con la misma forma de trie; la diferencia es la contabilidad que mantienen para APIs de uso general que no necesitamos.

El trie Darts de la Optimización 1 necesita dos cargas dependientes por paso de byte. Primero el desplazamiento base del padre, luego la entrada de verificación en el espacio del hijo candidato para confirmar que el padre lo posee. La entrada de verificación solo valida la propiedad. La misma pregunta puede ser respondida por un mapa de bits por nodo que registra qué bytes tienen hijos. Por la construcción del trie de doble matriz, el bit para el byte b en el mapa de bits del padre se establece exactamente cuando el espacio candidato para ese byte es propiedad del padre. El mapa de bits y la matriz de verificación llevan la misma información, por lo que podemos descartar la matriz de verificación de la estructura de tiempo de ejecución por completo. El indexado de ruta crítica no está verificado porque la construcción del trie y la estructura del bucle ya garantizan los límites. El intercambio es el Espacio: a 64 bytes por nodo × ~780K espacios, el trie crece de ~9 MB a ~50 MB.

La búsqueda en mapa de bits se compila en una única prueba de bit contra una palabra de 64 bits, que se ejecuta en aproximadamente un ciclo. La instrucción POPCNT relacionada, que cuenta todos los bits establecidos en un registro, es lo que usan otros diseños de trie indexados por mapa de bits para calcular un desplazamiento de hijo desde el mapa de bits directamente. No lo necesitamos porque la construcción Darts ya almacena ese desplazamiento en el 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)
}

Mientras estamos en ello, empaquetamos los cuatro campos por nodo (mapa de bits, base, ID de token y puntaje) en una línea de caché de 64 bytes. Una visita a un solo nodo ahora extrae todo: el puntaje y el ID del token que necesitamos en los nodos terminales, más el mapa de bits y la base para el siguiente paso. La matriz de verificación solo se usa durante la construcción y se descarta del diseño en tiempo de ejecución. El indexado en la ruta crítica no está marcado porque la construcción del trie y la estructura del bucle ya garantizan los límites. El intercambio es el Espacio: a 64 bytes por nodo × ~780K espacios el trie crece de ~9 MB a ~50 MB.

Resultados en 514 tokens:

Motor	p50	p99	Instrucciones	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%
Línea 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 y empaquetado en línea	65 µs	77 µs	1.06M	1.8K	20.53%	15.01%

Resultados del tokenizer con 514 tokens tras añadir bitmap y empaquetado en línea.

El reloj de pared cae sistemáticamente otro 4.5% en 514 tokens y crece con la longitud de la entrada, alcanzando el 6-7% en 4k y 16k. Descartar la matriz de verificación convierte el paso del trie por byte de dos cargas de línea de caché en una, lo que supone 2.5 veces menos accesos a L2 por codificación (1.8K frente a los 4.6K de Darts simple). La tasa de fallos salta del 6% al 20%, pero la mayor parte de eso es el denominador: con menos de la mitad de los accesos, los fallos que tenemos ocupan una mayor parte.

Los datos del trie son aproximadamente 50 MB (780,000 nodos a 64 bytes cada uno). En un proceso Linux predeterminado, esa memoria está respaldada por páginas de 4 KB, por lo que el trie completo abarca unas 12,000 páginas virtuales distintas, mucho más allá de lo que puede contener el búfer de búsqueda de traducción de datos de primer nivel de la CPU (Sapphire Rapids tiene 96 entradas). Cada paso de Viterbi toca un nodo diferente, y durante una codificación de 512 tokens quemamos aproximadamente 9,000 ciclos en recorridos de tablas de páginas, alrededor del 3% del presupuesto por codificación.

La solución es respaldar el trie con páginas enormes de 2 MB en lugar de las de 4 KB. Los mismos 50 MB ahora abarcan 25 páginas, caben completamente en el Búfer de búsqueda de traducción, y los directorios de la tabla de páginas dejan de competir con el trie por la caché L2.

Linux tiene dos formas de respaldar la memoria con páginas de 2 MB. La oportunista es las páginas enormes transparentes (THP): su programa solicita al kernel que actualice una región de memoria de 4 KB a páginas de 2 MB, y un hilo de kernel en segundo plano intenta hacerlo cuando tiene tiempo y memoria física contigua disponible. La otra alternativa es asignar el trie directamente desde un grupo de páginas enormes pre-reservado que el kernel reserva al arrancar:

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 grupo debe existir en el host (vm.nr_hugepages configurado al arrancar). En nuestras instancias de producción, se reservan 10,561 páginas enormes; el trie consume 24 de ellas. Después del mmap, /proc/<pid>/smaps confirma la región: KernelPageSize: 2048 kB, Private_Hugetlb: 49152 kB.

La ganancia es del 3–12% en reloj de pared a través de tamaños de entrada, consistente y proporcional a la longitud del conjunto de trabajo:

Tokens	Sin páginas enormes	Con páginas enormes	Δ
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 del tokenizer con y sin páginas enormes en distintos tamaños de entrada.

La mayor ganancia es en 4k tokens, donde el conjunto de trabajo es lo suficientemente grande como para que el tráfico de la tabla de páginas compitiera activamente con los datos del trie por el ancho de banda de L2. Más allá de 4k, la ganancia se reduce porque los fallos de L3 dominan y los recorridos de la tabla de páginas se convierten en una parte menor del tiempo total.

Al reunir las tres optimizaciones, aquí está el motor final medido contra cualquier otro tokenizador RoBERTa que probamos. Junto con los binarios nativos C++ SentencePiece y C IREE, también comparamos las cajas de envoltorio Rust listas para usar que un usuario de producción buscaría, sentencepiece e iree-tokenizer. La caja sentencepiece mide una latencia 1.6-1.9 veces mayor que el binario nativo C++ en cada tamaño de entrada, con la sobrecarga proveniente del envoltorio que asigna una lista nueva de piezas de token, cada una manteniendo su propia copia del texto de la pieza, en cada codificación. La caja iree-tokenizer dimensiona su búfer interno para cada entrada, lo que evita una asignación constante por codificación que el binario IREE nativo paga independientemente del tamaño de la entrada, y eso importa más en entradas cortas donde la asignación dominaría de otro modo el trabajo de tokenización real. En longitudes de entrada de producción, el costo constante se amortiza y la sobrecarga por llamada del envoltorio al copiar tokens de vuelta a Rust eleva su latencia por encima de la del binario nativo.

Latencia p50 y p99 por codificación entre implementaciones de tokenizer y tamaños de entrada.

El motor final ejecuta el mismo Viterbi sobre el mismo vocabulario de 250,000 tokens, no asigna nada en la ruta crítica y produce una salida exacta de token contra la referencia. Cada paso de byte es ahora una única carga de línea de caché de 64 bytes, una prueba de mapa de bits y una suma de enteros, por debajo de una sonda HashMap, una llamada al asignador y varias ramas de casos especiales en la implementación de la línea base. Las instrucciones por codificación cayeron de 3.66M a 1.04M, una caída de 3.5 veces. El reloj de pared coincide con esa proporción en entradas cortas y se amplía aún más en entradas largas, donde las asignaciones por token de la referencia comienzan a desbordar L2 y L3.

En la práctica, este trabajo redujo la utilización de CPU en la capa de preprocesamiento de nuestra pila de inferencia entre 5 y 6 veces y recortó milisegundos de dos dígitos de la latencia del reranker.

La mayor parte de la reducción de latencia provino de cambios ligeros que eliminaron el trabajo desperdiciado: descartar asignaciones por codificación, reemplazar el trie de niños-HashMap con una matriz doble y eliminar la sobrecarga de API genérica. El uso del mapa de bits existente por nodo como verificación de validez y el respaldo del trie con páginas enormes de 2 MB representaron en conjunto un 10-20% adicional sobre el cambio mayor a un trie de doble matriz, dependiendo del tamaño de la entrada.

Las cajas de envoltorio Rust listas para usar alrededor de SentencePiece e IREE ilustran el mismo punto a la inversa. Sobre los tokenizadores nativos C y C++, el límite FFI y las asignaciones por llamada que introducen los envoltorios son visibles en los números de latencia. La fuente Rust para el codificador Viterbi final está disponible 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áginas enormes — 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

El algoritmo de Viterbi

Implementación de referencia

Una nota sobre el preprocesamiento

Dónde gasta su tiempo la implementación de referencia

Cómo medimos

Por qué SentencePiece e IREE son más rápidos: el trie de doble matriz

Línea base: Un puerto de cero asignación de la referencia

Optimización 1: Trie de doble matriz sin abstracciones adicionales

Optimización 2: Mapa de bits y empaquetado en línea

Optimización 3: Páginas gigantes para el trie

Comparación final

Pensamientos finales

Referencias

Citación