El equipo de investigación de Perplexity AI reimplementó su tokenizador Unigram desde cero en Rust y abrió el código en pplx-garden, su repositorio de tecnología de inferencia.
En longitudes de entrada de producción, el nuevo codificador reduce la latencia p50 aproximadamente 5 veces en comparación con la caja de tokenizadores Hugging Face, ~2 veces en comparación con SentencePiece (C++) y ~1,5 veces en comparación con el tokenizador de IREE (C), con cero asignaciones de almacenamiento dinámico en estado estacionario. En producción, redujo la utilización de CPU en la pila de inferencia de Perplexity entre 5 y 6 veces y redujo en milisegundos de dos dígitos la latencia del reranker.
Por qué la tokenización se convirtió en un cuello de botella
El costo de inferencia de LLM generalmente se enmarca en el trabajo de GPU: cachés KV, núcleos de atención, enrutamiento experto. Pero los modelos más pequeños, como los modelos integrados, los clasificadores y los reordenadores, cuentan una historia diferente. Estos modelos son de dos a tres órdenes de magnitud más pequeños que los transformadores de frontera.
Un reclasificador que obtiene cientos de documentos de candidatos por solicitud es un claro ejemplo. Con un modelo pequeño, el cálculo de la GPU suele finalizar en milisegundos de un solo dígito. Cada entrada todavía pasa primero por la tokenización del lado de la CPU. Cuando el tamaño de los lotes es grande, la tokenización se convierte en una fracción significativa de la latencia total de las solicitudes.
El trabajo de Perplexity tiene como objetivo XLM-RoBERTa, un modelo con un vocabulario Unigram de 250.000 tokens entrenado con SentencePieza. Los codificadores de la familia RoBERTa ajustados son una opción de producción común para tareas de clasificación, recuperación y similitud.
¿Qué es la tokenización de Unigram?
Kudo introdujo la tokenización de Unigram en 2018 y se implementa en SentencePieza. Enmarca la segmentación como un problema de ruta más probable. Cada ficha de vocabulario tiene una probabilidad logarítmica aprendida. El tokenizador elige la segmentación cuyas puntuaciones de token suman el valor más alto.
El algoritmo utilizado para encontrar el mejor camino es el algoritmo de Viterbi, una técnica de programación dinámica de 1967. Las posiciones de los bytes forman capas de gráficos y los tokens de vocabulario son bordes que abarcan un rango de bytes contiguos. La recurrencia de DP itera sobre las posiciones de los bytes y actualiza la ruta de mejor puntuación en cada posición.
El bucle exterior se ejecuta en tiempo lineal en relación con la longitud de entrada. El bucle interno recorre un ensayo de vocabulario (una estructura de árbol de prefijos) en cada posición de byte. Con una entrada de 16.000 tokens, este paseo interior ejecuta cientos de miles de transiciones de prueba. Es el camino caliente.
¿Qué fue lento en la implementación de la cara de abrazo?
La caja de tokenizadores Hugging Face es el tokenizador de Rust predeterminado que utilizan la mayoría de los equipos. Perplejidad lo utilizó como referencia. Con 514 tokens (512 + inyección BOS/EOS), la implementación de referencia tenía tres patrones costosos:
La asignación por token es constante: aproximadamente 2 KB y ~18 asignaciones por token, independientemente del tamaño de entrada. El problema de latencia se vuelve grave en entradas más largas cuando las asignaciones acumulativas desbordan la caché L2 por núcleo.
Establecer una línea de base antes de cambiar el Trie
Antes de cambiar la estructura trie, Perplexity primero aisló cuánto costo provenía únicamente del trabajo innecesario. Hicieron un puerto de asignación cero de la referencia: el mismo intento de HashMap, pero con una estructura temporal propiedad de la persona que llama reutilizada en las llamadas y los ID de token almacenados directamente en los nodos de prueba (eliminando la asignación de cadenas por coincidencia y la búsqueda de mapa hash secundario).
Esta línea de base ya redujo la latencia de p50 a 155 µs en 514 tokens, frente a los 326 µs de la referencia. Las instrucciones retiradas cayeron 2,4x. El costo restante fue la propia búsqueda del puntero HashMap, que se abordó en el siguiente paso.
Las tres optimizaciones
Optimización 1: Trie de doble matriz
El experimento Hugging Face almacena a los niños en un HashMap en cada nodo. Cada paso de byte requiere un cálculo hash, dos desreferencias de puntero y un acceso al montón. Perplexity reemplazó esto con un trie de doble matriz, la misma estructura utilizada por SentencePiece e IREE, introducida originalmente por Aoe en 1989.
Un trie de doble matriz codifica el trie completo en dos matrices de enteros planos, base y check. Una búsqueda secundaria es: siguiente = base[nodo] + byte, luego verificar verificar[siguiente] == nodo. Es decir, dos lecturas de matriz, una suma de enteros y una comparación, sin hash ni persecución de punteros. Para el vocabulario de 250K de XLM-RoBERTa, todo el conjunto cabe en ~9 MB de memoria contigua. El conjunto de trabajo en caliente por codificación es del orden de 100 KB, que cabe en la caché L2.
A diferencia de SentencePiece e IREE, que son bibliotecas de uso general con contabilidad de celosía y canalizaciones de múltiples etapas, Perplexity incorporó el trie directamente en el bucle de Viterbi y eliminó esa sobrecarga por completo.
Resultado en 514 tokens: p50 cayó de 155 µs (línea de base de asignación cero) a 68 µs. El reloj de pared cayó 4,8 veces respecto a la referencia original.
Optimización 2: mapa de bits y empaquetado en línea
La prueba de doble matriz todavía requiere dos cargas de matriz dependientes por paso de byte: primero el desplazamiento base del padre, luego la matriz de verificación para confirmar que la transición es válida. Perplexity reemplazó la matriz de verificación con un mapa de bits por nodo (cuatro palabras de 64 bits, 32 bytes) que registra cuáles de los 256 bytes posibles tienen transiciones secundarias válidas.
Una búsqueda de mapa de bits se compila en una prueba de un solo bit con una palabra de 64 bits. La matriz de verificación se usa solo durante la construcción de prueba y se elimina por completo del diseño de tiempo de ejecución.
También empaquetaron los cuatro campos por nodo (mapa de bits, base, ID de token y puntuación) en una única línea de caché de 64 bytes, coincidiendo exactamente con el ancho de la línea de caché de la CPU. Un paso de prueba ahora carga una única línea de caché que cubre el mapa de bits para la verificación del siguiente byte, el desplazamiento base para la ranura secundaria y el ID del token y la puntuación en los nodos terminales.
Compensación: el tamaño del trie aumenta de ~9 MB a ~50 MB (780 000 nodos x 64 bytes). El conjunto de trabajo en caliente por codificación sigue siendo ~100 KB.
Resultado en 514 tokens: Reducción adicional del 4,5% del reloj de pared. Los accesos L2 cayeron de 4,6K a 1,8K por codificación.
Optimización 3: páginas enormes para Trie
Con 50 MB, el trie abarca aproximadamente 12.000 páginas virtuales en un sistema Linux predeterminado que utiliza páginas de 4 KB. El TLB de datos de primer nivel en Intel Sapphire Rapids tiene 96 entradas. Cada paso de Viterbi toca un nodo trie diferente, por lo que los errores de TLB se acumulan. En una codificación de 512 tokens, Perplexity estimó aproximadamente 9000 ciclos gastados en recorridos por tablas de páginas, aproximadamente el 3% del presupuesto por codificación.
Perplexity respaldó el trie con páginas enormes de 2 MB a través de mmap con la bandera MAP_HUGETLB. Los mismos 50 MB ahora abarcan 25 páginas, dentro del TLB. Esto requiere que vm.nr_hugepages esté configurado en el arranque. En producción se reservan 10.561 páginas enormes; el trie usa 24.
Resultado: Reducción del 3-12% del reloj de pared dependiendo de la longitud de entrada. La mayor ganancia se produjo en 4.098 tokens (-12,0%), donde el tráfico de la tabla de páginas competía activamente con los datos de prueba por el ancho de banda L2. Más allá de los tokens 4K, la ganancia se reduce porque dominan los fallos L3.
Resultados finales de referencia
Todas las mediciones son de un solo subproceso, fijadas a un núcleo en un Intel Xeon Platinum 8488C, con 10 000 iteraciones después de 1000 rondas de calentamiento. En 514 fichas:
En toda la secuencia de optimización, las instrucciones por codificación cayeron de 3,66 millones a 1,04 millones, una reducción de 3,5 veces. El reloj de pared iguala esa relación en entradas cortas y se amplía en entradas largas donde las asignaciones por token de la referencia desbordan L2 y L3.
Un hallazgo adicional: las cajas de envoltura Rust disponibles en el mercado alrededor de SentencePiece e IREE agregan una sobrecarga de latencia de 1,6 a 1,9 veces en comparación con los binarios nativos de C/C++. La caja de piezas de oración asigna una lista nueva de piezas simbólicas en cada llamada. Los gastos generales son mensurables pero se amortizan con insumos prolongados.
El codificador final de Perplexity produce una salida exacta del token frente a la referencia. En producción, utiliza rayón para paralelizar los núcleos.
Explicador visual de Marktechpost
Conclusiones clave
Perplexity reconstruyó su tokenizador Unigram apuntando al vocabulario SentencePiece de 250 000 tokens de XLM-RoBERTa. El nuevo codificador logra cero asignaciones de montón en estado estable y ~63 µs p50 en 514 tokens. Tres optimizaciones: trie de doble matriz, mapa de bits + empaquetamiento de línea de caché de 64 bytes y páginas grandes de 2 MB para el trie. Resultado intermedio: un puerto HashMap de asignación cero solo cortó p50 de 326 µs a 155 µs antes de que se cambiara el intento Impacto en la producción: reducción de 5 a 6 veces en la utilización de la CPU y reducción de ms de dos dígitos en la latencia del reranker
Consulte el repositorio y los detalles técnicos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 150.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros