Acelerando Sonar mediante especulación

La decodificación especulativa acelera la velocidad de generación de los Modelos de Lenguaje Extensos (LLM) utilizando un modelo 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 una ejecución del costoso modelo objetivo produzca un único token, se emiten múltiples 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 Sonar.

Decodificación especulativa

La decodificación especulativa aprovecha la estructura de los lenguajes naturales y la naturaleza autorregresiva de los transformers para acelerar la generación de tokens. Aunque los modelos más grandes, como Llama-70B, poseen más conocimiento que los más pequeños, como Llama-1B, en algunas tareas más simples su rendimiento es similar. Esta superposición sugiere que ciertas secuencias son generadas mejor por los modelos menos costosos, dejando los problemas complejos a los más grandes. El desafío radica en determinar qué finalizaciones son mejores y si la generación del modelo más pequeño tiene la misma calidad que la del más grande.

Afortunadamente, los LLM son transformers autorregresivos: cuando se les proporciona una secuencia de tokens, generan 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 permite la especulación: si un modelo más pequeño genera una secuencia de tokens a partir de un prefijo de entrada, esta puede ejecutarse a través del modelo más grande para determinar qué tan bien se alinea con el modelo objetivo. Cada prefijo de los candidatos se califica con una probabilidad y se elige el más largo que supere un umbral de aceptación. Como beneficio adicional, el modelo objetivo también proporciona un token posterior de forma gratuita: si un modelo borrador genera n tokens, se pueden emitir hasta n + 1 en un solo paso.

Ampliar

En el momento de la inferencia, el proceso de muestreo especulativo se puede dividir aproximadamente en 4 etapas:

Prefill: tanto el modelo objetivo como el modelo borrador deben ejecutarse en la secuencia de entrada para completar las entradas de la caché KV. Si bien algunos esquemas, como Medusa, utilizan capas densas más simples para la predicción, en esta publicación nos centramos en borradores basados en transformers que necesitan sus propias cachés KV.

Generación de borradores: el modelo borrador itera para producir una cantidad de tokens fijos. La secuencia de borrador puede ser lineal o el modelo puede explorar una estructura similar a un árbol hasta una profundidad determinada (EAGLE, Medusa). Aquí, nos centramos en secuencias lineales.

Aceptación: el modelo objetivo se ejecuta en la secuencia de borrador, creando logits correspondientes a cada token de borrador. Se determina la longitud de la secuencia aceptable más larga.

Generación del objetivo: dado que el objetivo generó logits, en la posición no coincidente o al final de la secuencia, los logits corresponden a un token posterior. Estos logits se pueden muestrear para proporcionar un token robusto del objetivo, finalizando la secuencia.

Existen varios métodos para implementar la decodificación especulativa. En esta publicación, nos centraremos en los esquemas que utilizamos para acelerar los modelos Sonar utilizando un modelo 1B interno, así como en los mecanismos de predicción que estamos desarrollando para acelerar modelos a la escala de DeepSeek.

Objetivo-Borrador

La decodificación especulativa se puede lograr acoplando un LLM pequeño existente como modelo borrador a un modelo objetivo para generar secuencias candidatas. En producción, hemos acelerado Sonar utilizando un modelo Llama-1B ajustado en el mismo conjunto de datos que el objetivo. Si bien este enfoque no requirió entrenar un borrador desde cero, el modelo pequeño aún utiliza una capacidad significativa de caché KV e introduce una ligera sobrecarga de prefill, lo que aumenta el TTFT.

Bajo este esquema, el decodificador solo especula sobre lotes de solo decodificación, generando tokens a través de muestreo estándar durante el prefill o en lotes mixtos de prefill-decodificación. En la etapa de prefill, los logits objetivo se muestrean inmediatamente para también prellenar el token recién generado en la caché KV del borrador. El borrador aún no se muestrea, pero los logits que produce se transfieren a la etapa de decodificación.

Ampliar

En la decodificación, el modelo borrador avanza, muestreando el token superior en cada etapa. Una vez alcanzada la longitud de borrador deseada, los tokens se ejecutan a través del modelo objetivo para producir los logits en función de los cuales el muestreador identifica la longitud de la secuencia aceptada. La aceptación se determina comparando las distribuciones de probabilidad completas del borrador y el objetivo. Dado que el objetivo siempre genera un conjunto de logits después de la secuencia de borrador aceptada, eso se muestrea para producir una salida adicional. Como el modelo borrador aún no ha visto ese token aceptado, se vuelve a ejecutar para completar sus entradas de caché KV correspondientes en preparación para el siguiente paso de decodificación, transfiriendo los logits nuevamente.

EAGLE

EAGLE es un esquema de decodificación especulativa que explora múltiples secuencias de borrador, generadas a través de un recorrido similar a un árbol de tokens de borrador probables. Se explora un árbol fijo (EAGLE) o de forma dinámica (EAGLE-2) utilizando ejecuciones consecutivas de los tokens de borrador, considerando los candidatos Top-K en cada nodo en lugar de seguir el token con mayor puntuación en una secuencia lineal. Luego, las secuencias se califican y se selecciona la más larga adecuada para continuar, añadiendo también un token adicional del objetivo.

Ampliar

Para lograr una predicción más precisa, un modelo borrador EAGLE predice no solo basándose en tokens, sino también utilizando las características objetivo (estados ocultos de la última capa) del modelo objetivo. La desventaja de EAGLE es la necesidad de entrenar modelos borradores pequeños y personalizados que sean lo suficientemente precisos como para generar candidatos adecuados dentro de un presupuesto de baja latencia. Normalmente, un modelo borrador es una sola capa de transformer idéntica a una capa de decodificador del modelo original, la cual está estrechamente acoplada al objetivo mediante la vinculación a sus proyecciones de embeddings y lm_head. Como esto requiere menos capacidad de caché KV, EAGLE tiene una huella de memoria menor.

Para verificar secuencias similares a un árbol en el modelo objetivo, se deben utilizar máscaras de atención personalizadas. Desafortunadamente, el uso de una máscara de atención personalizada para una secuencia completa ralentiza significativamente la atención para longitudes de entrada realistas (hasta un 50%), anulando 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 del árbol en producción, centrándonos en cambio en el caso especial de predicción de un solo token a través de esquemas tipo 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 tanto en las etapas de prefill como de decodificación en comparación con la especulación regular de borrador-objetivo. El modelo borrador utiliza tanto tokens como estados ocultos: el token t_{i+1} se muestrea a partir 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 buffers de tokens de entrada deben desplazarse un paso a la izquierda en relación con los vectores de estado oculto emitidos por el objetivo. La figura a continuación marca las correspondencias utilizadas para el entrenamiento, así como el desplazamiento durante la inferencia.

Ampliar

El flujo de decodificación es bastante similar a la decodificación borrador-objetivo, con la excepción de que se transfieren tanto los estados ocultos como los logits. Nuestra implementación comparte todos los logits de muestreo y procesamiento asociados, especializando solo las invocaciones de avance del modelo. Cuando se predicen múltiples tokens, el modelo borrador utiliza estados ocultos de borrador para la predicción, completando también las entradas de caché KV en función de sus propias características. A largo plazo, esto puede degradar la precisión. Posteriormente, al ejecutar el modelo borrador para completar la entrada de caché KV para la predicción objetivo, lo ejecutamos en toda la secuencia tomando los estados ocultos objetivo más precisos como entradas. Dado que estos modelos borradores son pequeños, el costo adicional de procesar los tokens adicionales es insignificante.

Entrenamiento de cabezales MTP

Para beneficiarnos de MTP, construimos la infraestructura necesaria para entrenar cabezales MTP conectados a nuestros modelos ajustados en los conjuntos de datos de Perplexity, ejecutándose en un nodo con 8 dispositivos H100. En aproximadamente un día, podemos construir cabezales para modelos que van desde Llama-1B hasta Llama-70B y DeepSeek V2-Lite. Para modelos más grandes, confiamos en los cabezales MTP construidos durante el proceso de ajuste.

El objetivo del entrenamiento de MTP es hacer coincidir los estados ocultos y logits del borrador extrapolados de los estados ocultos del objetivo con los logits y estados ocultos del siguiente token del objetivo. Como la inferencia para los estados ocultos es costosa, los precalculamos 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 debidas a la cuantización u optimizaciones no obstaculicen los resultados, para la pérdida de validación y la estimación de precisión reutilizamos completamente la implementación de inferencia tanto del modelo objetivo como del modelo borrador.

Al escalar desde el conjunto de datos ShareGPT utilizado en el artículo original a muestras más grandes, notamos que la arquitectura de cabezal MTP descrita e implementada en el artículo de EAGLE falló al entrenar 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 prompts sustancialmente más largos. Dado que los cabezales EAGLE originales divergían ligeramente en estructura de un transformer típico, reintrodujimos algunas capas de Normalización RMS que fueron eliminadas. Descubrimos que esto no solo permitió que el entrenamiento convergiera, sino que también aumentó la precisión de los cabezales en algunos puntos porcentuales.

Ampliar

Las normalizaciones de capa no solo facilitan el entrenamiento, sino que reintroducirlas también es matemáticamente intuitivo. Los cabezales MTP reutilizan los embeddings y las proyecciones de logits del modelo objetivo, ya que pueden ser de tamaño sustancial (alrededor de 2 GB para Llama 70B). Durante el entrenamiento, estos se congelan y la expectativa es que la capa MTP aprenda a integrar las predicciones en el mismo Espacio vectorial que lo que la capa de proyección del modelo original aprendió durante el entrenamiento. Al eliminar las normalizaciones, se espera que un solo MLP aprenda la misma función que un MLP seguido de una normalización, lo que obstaculiza la comparación entre los estados ocultos del borrador y los modelos 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 se deben asignar páginas en la caché KV para los siguientes tokens. Los tokens de entrada y la información de la página KV se empaquetan en un buffer transmitido a todos los rangos paralelos que ejecutan el modelo. Finalmente, los metadatos se copian en la memoria de la GPU y se ejecuta el modelo para producir los logits a partir de los cuales se muestrea el siguiente token.

A diferencia de ciertas implementaciones que acoplan libremente un servidor de inferencia de borrador y objetivo a través de un contenedor que orquesta las solicitudes entre ellos, nuestros pares borrador-objetivo están estrechamente acoplados y avanzan en la generación al unísono. La programación de 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 une un modelo con el servidor de inferencia general, ya que todos exponen la misma interfaz.

El tiempo de ejecución de inferencia en Perplexity se basa en FlashInfer, que determina los metadatos que deben crearse para configurar y programar el kernel de atención. Dadas algunas secuencias de entrada que forman un lote, para prefill, decodificación o verificación, se debe realizar trabajo del lado de la CPU para asignar buffers intermedios y completar ciertos buffers constantes utilizados en la atención. Este trabajo se suma al costo de la programación de lotes y la asignación de páginas KV, que también generan latencias que deben ocultarse para maximizar la utilización de la GPU.

Si bien paralelizamos completamente el trabajo del lado de la CPU y de la GPU para la inferencia sin especulación, descubrimos que el equilibrio CPU-GPU para la decodificación especulativa es más complejo. El desafío principal surge del hecho de que el número de tokens aceptados determina la longitud de la secuencia para una ejecución posterior, lo que introduce un punto de sincronización GPU-a-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 algo de margen para ocultar operaciones costosas de la CPU. Dado que los modelos más pequeños no se benefician del paralelismo de tensores, existe una discrepancia entre el número de rangos en los que se fragmentan un objetivo y un borrador. En nuestra implementación, el modelo borrador se ejecuta solo en el rango líder de un grupo TP.

Ampliar

Como se indicó anteriormente, un paso de decodificación transfiere logits a la siguiente ejecución. Esto nos permite superponer una ejecución del modelo borrador con el trabajo de programación de lotes del lado de la CPU. Una vez que se organiza el lote, las llamadas repetidas al muestreador y al borrador producen los tokens de borrador. En paralelo, el lote para la verificación se organiza para el modelo objetivo y se sincroniza con los trabajadores paralelos. Los logits objetivo se verifican y muestrean para determinar las longitudes de secuencia aceptadas. En este punto, es necesaria la sincronización GPU-a-CPU para determinar las longitudes de secuencia posteriores. Como el modelo borrador solo se ejecuta en el nodo líder, su lote se configura secuencialmente y su ejecución se inicia para completar sus entradas de caché KV con el token adicional que produjo el objetivo. Los logits producidos por esta ejecución del borrador en la ejecución actual se utilizarán para muestrear el primer token del borrador en la ejecución posterior. Lo más importante es que, mientras el borrador se está ejecutando, se puede programar el siguiente lote.

Programación MTP para un solo token

Si bien el tiempo de ejecución aún no proporciona exploración de árbol de borrador al estilo Eagle, implementamos un caso especial de este esquema, considerando una secuencia lineal de tokens de borrador producida por un modelo del tamaño de una sola capa de decodificador de transformer. Este esquema se puede utilizar para la predicción de borradores utilizando los pesos de código abierto de DeepSeek R1. El subcaso de predecir un solo token es interesante, ya que las capas MTP grandes logran tasas de aceptación lo suficientemente altas como para justificar su sobrecarga.

La programación MTP es algo más compleja, ya que el modelo borrador es mucho más rápido, ocultando menos latencia del lado de la CPU. Además, el borrador se fragmenta junto con el modelo objetivo, lo que requiere transferencias de memoria compartida para la información del lote. Una ejecución comienza transfiriendo información del lote y muestreando el primer token de los logits transferidos, de manera similar al esquema anterior. A continuación, el objetivo se ejecuta para validar los tokens, procesando 2 * D tokens, donde D es el tamaño del lote de decodificación. Esto es ideal para el micro-lote 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 transfieren a la siguiente ejecución del borrador, mientras que los logits se pasan al muestreador para su verificación.

Ampliar

Al realizar una cantidad limitada de trabajo adicional en la GPU, evitamos la sincronización CPU-a-GPU después de la aceptación de la secuencia de borrador. Después de que se desplazan los tokens de entrada de los objetivos, un kernel conecta los siguientes tokens objetivo en sus ubicaciones correspondientes. Luego, el borrador se vuelve a ejecutar con la misma información de lote que el objetivo, completando las entradas de caché KV y creando los logits y estados ocultos para la siguiente ejecución, realizando 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 borrador. En paralelo con la ejecución del borrador, las longitudes de la secuencia se determinan en la CPU y se inicia la programación del siguiente lote, sin tener que esperar a que finalice el trabajo de la GPU.

La sobrecarga del trabajo adicional en la capa de borrador no es perceptible en la atención; sin embargo, las capas MLP son más problemáticas. Dado que las instrucciones de multiplicación de matrices se rellenan hasta un límite de 64 a lo largo de la dimensión del número de tokens, si la duplicación no requiere significativamente más bloques, la sobrecarga queda oculta. Para secuencias de borrador más largas, la sobrecarga es más costosa y el esquema utilizado para modelos regulares de borrador-objetivo funciona mejor.

Referencias

Inferencia rápida de Transformers mediante 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 borrador dinámicos

EAGLE-3: Escalado de la aceleración de inferencia de Modelos de Lenguaje Extensos mediante pruebas en tiempo de entrenamiento

Medusa: Marco simple de aceleración de inferencia de LLM con múltiples cabezales de decodificación

Normalización de capa

FlashInfer: Motor de atención eficiente y personalizable para el servicio de inferencia de LLM

Informe técnico de DeepSeek-V3

Citación

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}