Aceleración de Sonar mediante especulación

La decodificación especulativa acelera la velocidad de generación de los Modelos de Lenguaje Grandes (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 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 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, contienen 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 se generan mejor mediante 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 modelo más grande.

Afortunadamente, los LLM son transformers autorregresivos: cuando se les proporciona 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 permite la especulación: si una secuencia de tokens es generada por un modelo más pequeño a partir de un prefijo de entrada, se puede ejecutar 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 selecciona 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 borrador deben ejecutarse en la secuencia de entrada para poblar 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 enfocamos en borradores basados en transformers que necesitan sus propias cachés KV.

Generación del borrador: el modelo borrador itera para producir una serie de tokens fijos. La secuencia del borrador puede ser lineal o el modelo puede explorar una estructura similar a un árbol hasta una profundidad determinada (EAGLE, Medusa). Aquí nos enfocamos 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 del objetivo: dado que el modelo 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 modelo objetivo, coronando la secuencia.

Existen varios métodos para implementar la decodificación especulativa. En esta publicación nos enfocaremos 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 con precisión (fine-tuned) 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 la caché KV e introduce una ligera sobrecarga de prefill, lo que aumenta el TTFT.

Bajo este esquema, el decodificador solo especula en lotes de solo decodificación, generando tokens mediante muestreo estándar durante el prefill o en lotes mixtos de prefill y decodificación. En la etapa de prefill, los logits del objetivo se muestrean inmediatamente para rellenar también el token recién generado en la caché KV del modelo 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 del objetivo. Dado que el modelo objetivo siempre emite un conjunto de logits después de la secuencia de borrador aceptada, este se muestrea para producir una salida adicional. Como el modelo borrador aún no ha visto ese token aceptado, se vuelve a ejecutar para poblar sus entradas correspondientes en la caché KV en preparación para el siguiente paso de decodificación, transfiriendo nuevamente los logits.

EAGLE

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

Ampliar

Para lograr una predicción más precisa, un modelo borrador EAGLE predice no solo basándose en los 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 personalizados y pequeños que sean lo suficientemente precisos como para generar candidatos adecuados dentro de un presupuesto de baja latencia. Por lo general, 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 su vinculación a sus embeddings y proyecciones de lm_head. Dado que esto requiere menos capacidad de caché KV, EAGLE tiene una huella de memoria menor.

Para verificar secuencias similares a árboles 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 toda una secuencia ralentiza significativamente la atención para longitudes de entrada realistas (hasta en un 50%), anulando parte de la aceleración que se puede lograr mediante la especulación. Por esta razón, aún no hemos implementado la exploración completa de árboles en producción, enfocándonos en su lugar en el caso especial de la 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 de borrador a objetivo, con la excepción de que se utilizan estados ocultos junto con los tokens para la predicción. Se debe realizar un poco más de trabajo tanto en la etapa de prefill como en la de decodificación en comparación con la especulación regular de borrador a 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, los cuales a su vez se derivan de los estados ocultos H_i. En consecuencia, los búferes de tokens de entrada deben desplazarse un paso hacia la izquierda en relación con los vectores de estados ocultos emitidos por el modelo objetivo. La siguiente figura marca las correspondencias utilizadas para el entrenamiento, así como el desplazamiento durante la inferencia.

Ampliar

El flujo de decodificación es muy similar al de la decodificación de borrador a objetivo, con la excepción de que tanto los estados ocultos como los logits se transfieren. Nuestra implementación comparte todo el muestreo asociado y el procesamiento de logits, especializando únicamente las invocaciones hacia adelante (forward) del modelo. Cuando se predicen múltiples tokens, el modelo borrador utiliza estados ocultos de borrador para la predicción, poblando también las entradas de la 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 poblar la entrada de la caché KV para la predicción del objetivo, lo ejecutamos en toda la secuencia tomando los estados ocultos del objetivo, que son 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 adjuntos a nuestros modelos ajustados con precisión 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 con precisión.

El objetivo del entrenamiento de MTP es hacer coincidir los estados ocultos del borrador y los logits extrapolados de los estados ocultos del objetivo con los logits del siguiente token y los estados ocultos del modelo objetivo. Dado que 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 MTP de inferencia 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 por completo la implementación de inferencia tanto del modelo objetivo como del borrador.

Al escalar desde el conjunto de datos ShareGPT utilizado en el documento original a muestras más grandes, notamos que la arquitectura de cabezales MTP descrita e implementada en el documento de EAGLE no logró entrenarse para modelos del tamaño de 70B. A diferencia de ShareGPT, que contenía una mayor cantidad de secuencias más cortas, entrenamos con una cantidad ligeramente menor de solicitudes (prompts) sustancialmente más largas. Dado que los cabezales originales de EAGLE diferían ligeramente en su estructura de un transformer típico, reintrodujimos algunas capas de normalización RMS que se habían eliminado. Descubrimos que esto no solo permitió que el entrenamiento convergiera, sino que también aumentó la precisión de los cabezales en unos pocos puntos porcentuales.

Ampliar

Las normas de capa no solo facilitan el entrenamiento, sino que su reintroducción también es matemáticamente intuitiva. Los cabezales MTP reutilizan los embeddings y las proyecciones de logits del modelo objetivo, ya que pueden ser de un tamaño considerable (aproximadamente 2 GB para Llama 70B). Durante el entrenamiento, estos se congelan y la expectativa es que la capa MTP aprenda a incrustar las predicciones en el mismo espacio vectorial que el que aprendió la capa de proyección del modelo original durante el entrenamiento. Al eliminar las normas, se espera que un único perceptrón multicapa (MLP) aprenda la misma función que un MLP seguido de una norma, lo que dificulta la coincidencia entre los estados ocultos del modelo borrador y los del objetivo.

Inferencia con decodificación especulativa

En el motor de inferencia, para generar tokens para las secuencias de entrada, primero deben agruparse en lotes de un tamaño razonable, y luego se deben asignar páginas en la caché KV para los siguientes tokens. Los tokens de entrada y la información de las páginas de la caché KV se empaquetan en un búfer que se transmite (broadcast) 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 de manera flexible un servidor de inferencia borrador y objetivo mediante un contenedor (wrapper) que orquesta las solicitudes entre ellos, nuestros pares de borrador y objetivo están estrechamente acoplados y avanzan en la generación al unísono. La programación de lotes (batch scheduling) y la asignación de páginas de la caché 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 exponen la misma interfaz.

El tiempo de ejecución de inferencia en Perplexity está diseñado en torno a FlashInfer, el cual determina los metadatos que deben construirse para configurar y programar el núcleo de atención (attention kernel). Dadas algunas secuencias de entrada que forman un lote, ya sea para prefill, decodificación o verificación, se debe realizar trabajo en el lado de la CPU para asignar búferes intermedios y rellenar ciertos búferes 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 de la caché KV, las cuales también incurren en latencias que deben ocultarse para maximizar la utilización de la GPU.

Si bien paralelizamos por completo el trabajo en el lado de la CPU y de la GPU para la inferencia sin especulación, descubrimos que el equilibrio entre CPU y GPU para la decodificación especulativa es más complejo. El desafío principal surge del hecho de que la cantidad de tokens aceptados determina la longitud de la secuencia para una ejecución posterior, lo que introduce un punto de sincronización entre GPU y CPU que es difícil de evitar. Experimentamos con diferentes esquemas de programación para ocultar de la melhor manera la latencia del trabajo de la CPU.

Programación de borrador a objetivo

A pesar de ser más pequeño que un modelo objetivo, cuando se utiliza un LLM completo como borrador, este sigue introduciendo una latencia considerable en la GPU, lo que proporciona cierto margen para ocultar operaciones costosas en la CPU. Dado que los modelos más pequeños no se benefician del paralelismo tensorial, existe una discrepancia entre el número de rangos en los que se divide (sharded) un modelo objetivo y uno borrador. En nuestra implementación, el modelo borrador se ejecuta únicamente en el rango líder de un grupo de paralelismo tensorial (TP).

Ampliar

Como se indicó anteriormente, un paso de decodificación transfiere los logits a la siguiente ejecución. Esto nos permite solapar una ejecución del modelo borrador con el trabajo de programación de lotes en el lado de la CPU. Una vez armado el lote, las llamadas repetidas al muestreador y al borrador producen los tokens del borrador. En paralelo, se arma el lote para la verificación correspondiente al modelo objetivo y se sincroniza con los trabajadores paralelos. Los logits del objetivo se verifican y se muestrean para determinar las longitudes de las secuencias aceptadas. En este punto, es necesaria la sincronización entre GPU y CPU para determinar las longitudes de las secuencias posteriores. Dado que el modelo borrador solo se ejecuta en el nodo líder, su lote se configura de manera secuencial y se inicia su ejecución para poblar sus entradas de la caché KV con el token adicional que produjo el modelo 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 la exploración de árboles de borrador al estilo de 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 los cabezales 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 borrador es mucho más rápido y oculta menos latencia en el lado de la CPU. Además, el borrador se divide (sharded) junto con el modelo objetivo, lo que requiere transferencias de memoria compartida para la información de los lotes. Una ejecución comienza transfiriendo la información del lote y muestreando el primer token a partir de los logits transferidos, de manera similar al esquema anterior. A continuación, se ejecuta el modelo objetivo para validar los tokens, procesando 2 * D tokens, donde D es el tamaño del lote de decodificación. Esto es ideal para lotes pequeños (micro-batching) en modelos de mezcla de expertos (MoE) a través de interconexiones más lentas como InfiniBand, ya que el lote se divide equitativamente 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 entre CPU y GPU después de la aceptación de la secuencia del borrador. Una vez que se desplazan los tokens de entrada de los objetivos, un núcleo (kernel) inserta los siguientes tokens del objetivo en sus ubicaciones correspondientes. A continuación, el borrador se vuelve a ejecutar con la misma información de lote que el objetivo, poblando las entradas de la caché KV y construyendo 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 las secuencias se determinan en la CPU y se inicia la programación del siguiente lote, sin tener que esperar a que finalice el trabajo en la GPU.

La sobrecarga de 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 (pad) 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 los modelos regulares de borrador a objetivo funciona mejor.

Referencias

Inferencia rápida a partir 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 Grandes mediante pruebas en tiempo de entrenamiento

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

Normalización de capas

FlashInfer: Motor de atención eficiente y personalizable para la prestación 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}
}