Artículo
Acelerando Sonar a Través de la Especulación
La decodificación especulativa acelera la velocidad de generación de los Modelos de Lenguaje Grandes (LLMs) al usar un modelo de borrador rápido y pequeño para producir candidatos de finalización que son verificados por el modelo objetivo más grande.
Bajo este esquema, en lugar de que el costoso objetivo produzca un solo token, se emiten varios en un solo paso. Aquí presentamos los detalles de implementación de varios tipos de decodificación especulativa, aplicados en Perplexity para reducir la latencia entre tokens en los modelos de Sonar.
Decodificación Especulativa
La Decodificación Especulativa aprovecha la estructura de los lenguajes naturales y la naturaleza autorregresiva de los transformadores para acelerar la generación de tokens. Aun cuando los modelos más grandes, como Llama-70B, contienen más conocimiento que los más pequeños, como Llama-1B, en algunas tareas más simples actúan de manera similar. Esta superposición sugiere que ciertas secuencias son mejor generadas por los modelos menos costosos, dejando los problemas complejos a los más grandes. El desafío reside en determinar qué completaciones son mejores y si la generación del modelo más pequeño es de la misma calidad que la del más grande.
Afortunadamente, los LLMs son transformadores autorregresivos: al recibir una secuencia de tokens, emiten la distribución de probabilidad del siguiente token. Además, los logits derivados de las características intermedias asociadas con los tokens en la secuencia de entrada también indican qué tan probable es que el modelo emita esos tokens exactos. Esta propiedad facilita la especulación: si una secuencia de tokens es generada por uno más pequeño partiendo de un prefijo de entrada, puede ejecutarse a través del más grande para determinar qué tan bien se alinea con el modelo objetivo. Cada prefijo de los candidatos se puntúa con una probabilidad y se elige el más largo por encima de un umbral de aceptación. Como ventaja, el modelo objetivo también proporciona un token posterior de manera gratuita: si un modelo de borrador genera n tokens, se pueden emitir hasta n + 1 en un paso.
En el tiempo de inferencia, el proceso de muestreo especulativo puede dividirse en aproximadamente 4 etapas:
Prefill: tanto los modelos objetivo como de borrador deben ejecutarse en la secuencia de entrada para llenar las entradas de la memoria caché KV. Mientras que algunos esquemas, como Medusa, usan capas densas más simples para la predicción, en esta publicación nos centramos en borradores basados en transformadores que necesitan sus propias memorias caché KV.
Generación de borrador: el modelo de borrador itera para producir un número fijo de tokens. La secuencia del borrador puede ser lineal o el modelo puede explorar una estructura de árbol hasta una profundidad dada (EAGLE, Medusa). Aquí, nos centramos en secuencias lineales.
Aceptación: el modelo objetivo se ejecuta en la secuencia del borrador, construyendo logits correspondientes a cada token del borrador. Se determina la longitud de la secuencia aceptable más larga.
Generación objetivo: dado que el objetivo generó logits, en la posición desajustada o al final de la secuencia, los logits corresponden a un token posterior. Estos logits pueden muestrearse para proporcionar un token robusto del objetivo, completando la secuencia.
Existen varios métodos para implementar la decodificación especulativa. En esta publicación, nos centraremos en los esquemas que usamos para acelerar los modelos Sonar usando un modelo interno de 1B, así como los mecanismos de predicción que estamos desarrollando para acelerar modelos a la escala de DeepSeek.
Objetivo-Borrador
La decodificación especulativa puede lograrse acoplando un LLM pequeño existente como modelo de borrador a un modelo objetivo para generar secuencias candidatas. En producción, hemos acelerado Sonar usando un modelo Llama-1B ajustado con el mismo conjunto de datos que el objetivo. Si bien este enfoque no requería entrenar un borrador desde cero, el modelo pequeño todavía utiliza una capacidad significativa de memoria caché KV e introduce una leve sobrecarga de prellenado, aumentando el TTFT.
Bajo este esquema, el decodificador solo especula en lotes de solo decodificación, generando tokens a través de muestreo estándar durante el prellenado o en lotes mixtos de prellenado-decodificación. En la etapa de prellenado, los logits del objetivo se muestrean inmediatamente para también prellenar el token recién generado en la memoria caché KV del borrador. El borrador aún no se muestrea, pero los logits que produce se llevan a la etapa de decodificación.
En la decodificación, se avanza el modelo de borrador, muestreando el token superior en cada etapa. Después de alcanzar la longitud deseada del borrador, los tokens se ejecutan a través del modelo objetivo para producir los logits sobre cuya base el muestreador identifica la longitud de la secuencia aceptada. La aceptación se determina comparando las distribuciones de probabilidad completas del borrador y del objetivo. Dado que el objetivo siempre emite un conjunto de logits siguiendo la secuencia de borrador aceptada, se muestrea para producir una salida adicional. Dado que el modelo de borrador aún no ha visto ese token aceptado, se vuelve a ejecutar para llenar sus entradas de memoria caché KV correspondientes en preparación para el siguiente paso de decodificación, llevando los logits nuevamente.
EAGLE
EAGLE es un esquema de decodificación especulativa que explora múltiples secuencias de borrador, generadas a través de una travesía en forma de árbol de tokens de borrador probables. Se explora un árbol fijo (EAGLE) o moldeado dinámicamente (EAGLE-2) usando ejecuciones consecutivas de los tokens de borrador, considerando los candidatos Top-K en cada nodo en lugar de seguir el token de mayor puntuación en una secuencia lineal. Las secuencias se puntúan y se selecciona la más larga adecuada para continuar, también agregando un token adicional del objetivo.
Para lograr una predicción más precisa, un modelo de borrador EAGLE predice no solo en base a tokens, sino también utilizando las características objetivo (estados escondidos de la última capa) del modelo objetivo. La desventaja de EAGLE es la necesidad de entrenar modelos de borrador pequeños personalizados que sean lo suficientemente precisos para generar candidatos adecuados dentro de un presupuesto de baja latencia. Típicamente, un modelo de borrador es una capa de transformador única idéntica a una capa de decodificador del modelo original, que está estrechamente acoplada al objetivo al conectarse a sus proyecciones de incrustaciones y lm_head. Dado que esto requiere menos capacidad de memoria caché KV, EAGLE tiene una menor huella de memoria.
Para verificar secuencias en forma de árbol en el modelo objetivo, deben utilizarse máscaras de atención personalizadas. Desafortunadamente, usar una máscara de atención personalizada para toda una secuencia reduce significativamente la velocidad de atención para longitudes de entrada realistas (hasta en un 50%), lo que anula parte de la aceleración alcanzable a través de la especulación. Por esta razón, aún no hemos implementado la exploración completa de árboles en producción, enfocándonos en cambio en el caso especial de predicción de un solo token mediante esquemas similares a MTP presentados en el Informe Técnico de DeepSeek-V3.
MTP
Este esquema es similar a la decodificación borrador-objetivo, con la excepción de que se utilizan estados ocultos junto con tokens para la predicción. Se debe realizar un poco más de trabajo en las etapas tanto de prellenado como de decodificación en comparación con la especulación borrador-objetivo regular. El modelo de borrador utiliza tanto tokens como estados ocultos: el token t_{i+1} se muestrea de los logits L_i correspondientes al token t_i, que a su vez se derivan de los estados ocultos H_i. En consecuencia, los búferes de tokens de entrada deben desplazarse un paso a la izquierda en relación a los vectores de estado oculto de salida del objetivo. La figura a continuación marca las correspondencias utilizadas para el entrenamiento, así como el desplazamiento durante la inferencia.
El flujo de decodificación es bastante similar a la decodificación borrador-objetivo, con la excepción de que se acarrea tanto los estados ocultos como los logits. Nuestra implementación comparte todo el muestreo asociado y el procesamiento de logits, especializando solo las invocaciones hacia adelante del modelo. Cuando se predicen múltiples tokens, el modelo de borrador utiliza estados ocultos del borrador para la predicción, también llenando entradas de memoria caché KV basadas en sus propias características. A largo plazo, esto puede degradar la precisión. Posteriormente, al ejecutar el modelo de borrador para llenar la entrada de memoria caché KV del objetivo, lo ejecutamos en toda la secuencia tomando los estados ocultos del objetivo más precisos como insumos. Como estos modelos de borrador son pequeños, el costo adicional de procesar los tokens adicionales es despreciable.
Entrenamiento de Cabezas MTP
Para beneficiarse de MTP, construimos la infraestructura necesaria para entrenar cabezas MTP adjuntas a nuestros modelos ajustados con los datasets de Perplexity, ejecutándose en un nodo con dispositivos 8xH100. En aproximadamente un día, podemos construir cabezas para modelos que van desde Llama-1B hasta Llama-70B y DeepSeek V2-Lite. Para modelos más grandes, confiamos en cabezas MTP construidas durante el proceso de ajuste fino.
El objetivo del entrenamiento MTP es hacer coincidir los estados ocultos del borrador y los logits extrapolados de los estados ocultos del objetivo con los siguientes logits de tokens y estados ocultos del objetivo. Dado que la inferencia para los estados ocultos es costosa, los precomputamos utilizando nuestra implementación optimizada para inferencia del modelo objetivo, para ser utilizados durante el entrenamiento. Sin embargo, para validar la implementación de inferencia MTP y garantizar que las diferencias numéricas debido a la cuantización o las optimizaciones no obstaculicen los resultados, para la estimación de pérdidas de validación y precisión reutilizamos por completo la implementación de inferencia tanto de los modelos objetivo como del borrador.
Al escalar desde el dataset ShareGPT utilizado en el papel original a muestras más grandes, notamos que la arquitectura de cabezas MTP esbozada e implementada en el artículo de EAGLE fallaba en el entrenamiento para modelos de tamaño 70B. A diferencia de ShareGPT que contenía un mayor número de secuencias más cortas, entrenamos en un número ligeramente menor de entradas sustancialmente más largas. Dado que las cabezas EAGLE originales divergían ligeramente en estructura de un transformador típico, reintrodujimos algunas capas de Normalización RMS que habían sido eliminadas. Encontramos que esto no solo permitió que el entrenamiento convergiera, sino que también mejoró la precisión de las cabezas en unos pocos puntos porcentuales.
No solo las normas de capas facilitan el entrenamiento, reintroducir las normas también es matemáticamente intuitivo. Las cabezas MTP reutilizan las incrustaciones y las proyecciones de logits del modelo objetivo, ya que pueden ser sustanciales en tamaño (aproximadamente 2 GB para Llama 70B). Durante el entrenamiento, estas se congelan y la expectativa es que la capa MTP aprenda a incrustar predicciones en el mismo espacio vectorial que la capa de proyección del modelo original aprendió durante el entrenamiento. Al eliminar las normas, se espera que un único MLP aprenda la misma función que un MLP seguido de una norma, lo que obstaculiza el emparejamiento entre los estados ocultos de los modelos de borrador y objetivo.
Inferencia con Decodificación Especulativa
En el motor de inferencia, para generar tokens para secuencias de entrada, primero deben agruparse en lotes de tamaño razonable, luego deben asignarse páginas en la memoria caché KV para los siguientes tokens. Los tokens de entrada y la información de la página KV luego se empaquetan en un búfer transmitido a todas las filas paralelas ejecutando el modelo. Finalmente, los metadatos se copian en la memoria GPU y el modelo se ejecuta para producir los logits de los cuales se muestrea el siguiente token.
A diferencia de ciertas implementaciones que acoplan de manera suelta un servidor de inferencia de borrador y objetivo a través de un envoltorio que orquesta las solicitudes entre ellos, nuestros pares borrador-objetivo están estrechamente acoplados y avanzan a través de la generación al unísono. La programación por lotes y la asignación de páginas KV se comparten entre los modelos para todas las formas de decodificación especulativa: esto unifica la lógica que conecta un modelo con el servidor de inferencia general, ya que todos ellos exponen la misma interfaz.
El tiempo de ejecución de inferencia en Perplexity está modelado en torno a FlashInfer, que determina los metadatos que necesitan construirse para configurar y programar el kernel de atención. Dadas algunas secuencias de entrada que forman un lote, para prellenado, decodificación o verificación, se debe realizar trabajo en el lado de la CPU para asignar búferes intermedios y llenar ciertos búferes constantes utilizados en atención. Este trabajo es adicional al costo de la programación por lotes y la asignación de páginas KV, las cuales también incurren en latencias que deben ocultarse para maximizar la utilización de la GPU.
Si bien paralelizamos completamente los trabajos en el lado de la CPU y la GPU para inferencia sin especulación, encontramos que el equilibrio CPU-GPU para la decodificación especulativa es más intricado. El principal desafío surge del hecho de que el número de tokens aceptados determina la longitud de la secuencia para una ejecución subsiguiente, introduciendo un punto de sincronización GPU-CPU difícil de evitar. Experimentamos con diferentes esquemas de programación para ocultar mejor la latencia del trabajo de la CPU.
Programación Borrador-Objetivo
A pesar de ser más pequeño que un modelo objetivo, cuando se utiliza un LLM completo como borrador, todavía introduce una latencia considerable en la GPU, proporcionando cierto margen para ocultar operaciones de CPU costosas. Como los modelos más pequeños no se benefician del paralelismo de tensores, hay una disonancia en el número de filas en las que un objetivo y un borrador se dividen. En nuestra implementación, el modelo de borrador se ejecuta solo en la fila líder de un grupo TP.
Como se indicó antes, un paso de decodificación acarrea logits al siguiente ciclo. Esto nos permite solapar una ejecución del modelo de borrador con el trabajo de programación de lotes en el lado de la CPU. Una vez que el lote se ha creado, las llamadas repetidas al muestreador y al borrador producen los tokens de borrador. En paralelo, el lote para verificación se reúne para el modelo objetivo y se sincroniza con los trabajadores en paralelo. Los logits objetivos son verificados y muestreados para determinar las longitudes de secuencias aceptadas. En este punto, la sincronización GPU-CPU es necesaria para determinar las longitudes de las secuencias subsiguientes. Dado que el modelo de borrador solo se ejecuta en el nodo líder, su lote se organiza secuencialmente y se inicia su ejecución para llenar sus entradas de memoria caché KV con el token adicional que produjo el objetivo. Los logits producidos por esta ejecución de borrador en la ejecución anterior se usarán para muestrear el primer token de borrador en el siguiente ciclo. Lo más importante, mientras el borrador esté en ejecución, el próximo lote puede programarse.
Programación MTP para un Solo Token
Si bien el tiempo de ejecución aún no proporciona una exploración de árboles de borrador al estilo Eagle, implementamos un caso especial de este esquema, considerando una secuencia lineal de tokens de borrador producidos por un modelo del tamaño de una sola capa de decodificador de transformador. Este esquema puede usarse para la predicción de borradores utilizando los pesos de código abierto de DeepSeek R1. El sub-caso de predecir un solo token es interesante, ya que las capas MTP grandes logran tasas de aceptación suficientemente altas como para justificar su sobrecarga.
La programación MTP es algo más compleja, ya que el modelo de borrador es mucho más rápido, ocultando menos latencia del lado de la CPU. Además, el borrador se divide junto al modelo objetivo, requiriendo transferencias de memoria compartida para la información por lotes. Un ciclo comienza transfiriendo la información por lotes y muestreando el primer token de los logits acarreados, de manera similar al esquema anterior. A continuación, el objetivo se ejecuta para validar tokens, procesando 2 * D tokens, donde D es el tamaño del lote de decodificación. Esto es ideal para micro-batching en modelos de Mezcla de Expertos (MoE) sobre interconexiones más lentas como InfiniBand, ya que el lote se divide uniformemente en dos mitades. Los estados ocultos del objetivo se trasladan al siguiente ciclo de borrador, mientras que los logits se pasan al muestreador para verificación.
Al realizar una cantidad limitada de trabajo adicional en la GPU, evitamos la sincronización CPU-GPU después de la aceptación de la secuencia de borrador. Después de que los tokens de entrada de los objetivos son desplazados, un kernel conecta los siguientes tokens objetivo en sus ubicaciones correspondientes. Luego, el borrador se vuelve a ejecutar con la misma información por lotes que el objetivo, llenando las entradas de caché KV y construyendo los logits y estados ocultos para la siguiente ejecución, haciendo algo de trabajo redundante en los tokens que no fueron aceptados. En estas situaciones, la latencia del trabajo no utilizado es apenas medible debido al pequeño tamaño del modelo de borrador. En paralelo con la ejecución del borrador, se determinan las longitudes de las secuencias en la CPU y se inicia la programación del siguiente lote, sin necesidad de esperar a que el trabajo de la GPU termine.
La sobrecarga de trabajo adicional en la capa de borrador no es observable en la atención, sin embargo, las capas MLP son más problemáticas. Dado que las instrucciones de multiplicación de matrices se ajustan a un límite de 64 a lo largo de la dimensión del número de tokens, si duplicar no requiere significativamente más bloques, la sobrecarga se oculta. Para secuencias de borrador más largas, la sobrecarga es más costosa y el esquema usado para modelos regulares borrador-objetivo funciona mejor.
Referencias
Inferencia Rápida de Transformadores a través de Decodificación Especulativa
EAGLE: El Muestreo Especulativo Requiere Repensar la Incertidumbre de las Características
EAGLE-2: Inferencia más Rápida de Modelos de Lenguaje con Árboles de Borradores Dinámicos
Medusa: Marco de Aceleración de Inferencia de LLM Simple con Cabezas de Decodificación Múltiples
FlashInfer: Motor de Atención Eficiente y Personalizable para Servicio de Inferencia de LLM