Aceleración de Sonar mediante especulación

A decodificación especulativa acelera a velocidade de xeración de grandes modelos de linguaxe (LLM) utilizando un modelo preliminar rápido e pequeno para producir candidatos de finalización que son verificados polo modelo obxectivo máis grande. Con este esquema, en lugar dunha execución do modelo obxectivo custoso que produce un único token, emítense varios nun só paso. Aquí presentamos os detalles de implementación de varios tipos de decodificación especulativa, aplicados en Perplexity para reducir a latencia entre tokens nos modelos Sonar.

Decodificación especulativa

A decodificación especulativa aproveita a estrutura das linguaxes naturais e a natureza autorregresiva dos transformadores para acelerar a xeración de tokens. Aínda que os modelos máis grandes, como Llama-70B, conteñen máis coñecemento que os máis pequenos, como Llama-1B, nalgunhas tarefas máis sinxelas funcionan de xeito similar. Esta superposición suxire que certas secuencias son xeradas mellor polos modelos menos custosos, deixando os problemas complexos aos máis grandes. O desafío radica en determinar que finalizacións son mellores e se a xeración do modelo máis pequeno ten a mesma calidade que a do máis grande.

Afortunadamente, os LLM son transformadores autorregresivos: cando se lles dá unha secuencia de tokens, emiten a distribución de probabilidade do seguinte token. Ademais, os logits derivados das características intermedias asociadas aos tokens da secuencia de entrada tamén indican a probabilidade de que o modelo emita eses tokens exactos. Esta propiedade permite a especulación: se unha secuencia de tokens é xerada por unha máis pequena a partir dun prefixo de entrada, pode executarse a través da máis grande para determinar como se aliña co modelo obxectivo. Cada prefixo dos candidatos valórase cunha probabilidade e elixido o máis longo por riba dun limiar de aceptación. Como bonificación, o modelo obxectivo tamén proporciona un token posterior de balde: se un modelo preliminar xera n tokens, pódense emitir ata n + 1 nun só paso.

Ampliar

No momento da inferencia, o proceso de a muestreo especulativo pódese dividir aproximadamente en 4 etapas:

Preenchemento: tanto o modelo obxectivo como o preliminar deben executarse na secuencia de entrada para encher as entradas da caché KV. Aínda que algúns esquemas, como Medusa, usan capas densas máis sinxelas para a predición, nesta publicación centramosnos en borradouros baseados en transformadores que necesitan as súas propias cachés KV.

Xeración preliminar: o modelo preliminar iteración tras iteración produce un número de tokens fixos. A secuencia preliminar pode ser lineal ou o modelo pode explorar unha estrutura en árbore ata unha profundidade dada (EAGLE, Medusa). Aquí, centramosnos en secuencias lineais.

Aceptación: o modelo obxectivo execútase na secuencia preliminar, construíndo logits correspondentes a cada token preliminar. Determínase a lonxitude da secuencia aceptable máis longa.

Xeración obxectivo: dado que o obxectivo xerou logits, na posición non coincidente ou na parte traseira da secuencia, os logits corresponden a un token posterior. Estes logits pódense mostraxe para proporcionar un token robusto do obxectivo, coroando a secuencia.

Existen varios métodos para implementar a decodificación especulativa. Nesta publicación, centrarémonos nos esquemas que utilizamos para acelerar os modelos Sonar utilizando un modelo 1B interno, así como nos mecanismos de predición que estamos a construír para acelerar os modelos á escala de DeepSeek.

Obxectivo-Preliminar

A decodificación especulativa pódese lograr acoplando un pequeno LLM existente como modelo preliminar a un modelo obxectivo para xerar secuencias candidatas. En produción, aceleramos Sonar utilizando un modelo Llama-1B axustado no mesmo conxunto de datos que o obxectivo. Aínda que este enfoque non requiriu adestrar un borrador desde cero, o modelo pequeno aínda usa unha capacidade significativa de caché KV e introduce unha lixeira sobrecarga de preenchemento, aumentando o TTFT.

Baixo este esquema, o decodificador só especula en lotes de só decodificación, xerando tokens mediante mostraxe estándar durante o preenchemento ou en lotes mixtos de preenchemento e decodificación. Na etapa de preenchemento, os logits obxectivo móstranse inmediatamente para preencher tamén o token recén xerado na caché KV do borrador. O borrador aínda non se mostra, pero os logits que produce transpórtanse á etapa de decodificación.

Ampliar

Na decodificación, o modelo preliminar avanza, mostrando o token superior en cada etapa. Unha vez alcanzada a lonxitude preliminar desexada, os tokens pasan polo modelo obxectivo para producir os logits en función dos cales o mostreador identifica a lonxitude da secuencia aceptada. A aceptación determínase comparando as distribucións de probabilidade completas do borrador e do obxectivo. Dado que o obxectivo sempre emite un conxunto de logits despois da secuencia preliminar aceptada, este móstrase para producir unha saída adicional. Dado que o modelo preliminar aínda non viu ese token aceptado, vólvese executar para encher as súas entradas de caché KV correspondentes en preparación para o seguinte paso de decodificación, transportando de novo os logits.

EAGLE

EAGLE é un esquema de decodificación especulativa que explora múltiples secuencias preliminares, xeradas a través dun percorrido en árbore de tokens preliminares probables. Explórase unha árbore fixa (EAGLE) ou de forma dinámica (EAGLE-2) mediante execucións consecutivas dos tokens preliminares, considerando os candidatos Top-K en cada nodo en lugar de seguir o token coa puntuación máis alta nunha secuencia lineal. As secuencias son entón puntuadas e selecciónase a máis longa axeitada para continuar, engadindo tamén un token adicional do obxectivo.

Ampliar

Para lograr unha predición máis precisa, un modelo preliminar EAGLE predí non só baseándose en tokens, senón tamén utilizando as características obxectivo (estados ocultos da última capa) do modelo obxectivo. A desvantaxe de EAGLE é a necesidade de adestrar modelos preliminares pequenos e personalizados que sexan o suficientemente precisos como para xerar candidatos axeitados dentro dun orzamento de baixa latencia. Normalmente, un modelo preliminar é unha única capa de transformador idéntica a unha capa de decodificador do modelo orixinal, que está fortemente acoplada ao obxectivo mediante a vinculación ás súas incrustacións e proxeccións lm_head. Dado que isto require menos capacidade de caché KV, EAGLE ten unha pegada de memoria menor.

Para verificar secuencias en forma de árbore no modelo obxectivo, débense usar máscaras de atención personalizadas. Desafortunadamente, usar unha máscara de atención personalizada para toda unha secuencia ralentiza significativamente a atención para lonxitudes de entrada realistas (ata nun 50 %), anulando parte da aceleración alcanzable mediante a especulación. Por este motivo, aínda non despregamos a exploración de árbores completas en produción, centrándonos no cambio no caso especial de predición dun único token mediante esquemas similares a MTP presentados no informe técnico de DeepSeek-V3.

MTP

Este esquema é similar á decodificación de borrador-obxectivo, coa excepción de que se usan estados ocultos xunto cos tokens para a predición. Debe facerse un pouco máis de traballo tanto na etapa de preenchemento como na de decodificación en comparación coa especulación habitual de borrador-obxectivo. O modelo preliminar usa tanto tokens como estados ocultos: o token t_{i+1} móstrase a partir dos logits L_i correspondentes ao token t_i, que á súa vez se derivan dos estados ocultos H_i. En consecuencia, os búferes de tokens de entrada deben desprazarse un paso á esquerda en relación cos vectores de estado oculto emitidos polo obxectivo. A figura seguinte marca as correspondencias utilizadas para o adestramento, así como o desprazamento durante a inferencia.

Ampliar

O fluxo de decodificación é moi similar ao da decodificación de borrador-obxectivo, coa excepción de que se transportan tanto os estados ocultos como os logits. A nosa implementación comparte todos os logits de mostraxe e procesamento de logits asociados, especializando só as invocacións cara adiante do modelo. Cando se predín varios tokens, o modelo preliminar utiliza estados ocultos preliminares para a predición, enchendo tamén entradas de caché KV baseadas nas súas propias características. A longo prazo, isto pode degradar a precisión. Posteriormente, ao executar o modelo preliminar para encher a entrada da caché KV para a predición obxectivo, executámolo en toda a secuencia tomando os estados ocultos obxectivo máis precisos como entradas. Dado que estes modelos preliminares son pequenos, o custo engadido de procesar os tokens adicionais é desprezable.

Adestramento de cabezas MTP

Para beneficiarnos de MTP, construímos a infraestrutura necesaria para adestrar cabezas MTP conectadas aos nosos modelos axustados nos conxuntos de datos de Perplexity, executándose nun nodo con 8 dispositivos H100. En aproximadamente un día, podemos construír cabezas para modelos que van desde Llama-1B ata Llama-70B e DeepSeek V2-Lite. Para modelos máis grandes, confiamos en cabezas MTP construídas durante o proceso de axuste fino.

O obxectivo do adestramento MTP é facer coincidir os estados ocultos preliminares e os logits extrapolados dos estados ocultos obxectivo cos logits do seguinte token e os estados ocultos do obxectivo. Dado que a inferencia para os estados ocultos é custosa, precomputámolos utilizando a nosa implementación optimizada para inferencia do modelo obxectivo, para ser utilizada durante o adestramento. Non entanto, para validar a implementación de inferencia MTP e garantir que as diferenzas numéricas debidas á cuantificación ou ás optimizacións non prexudiquen os resultados, para a perda de validación e a estimación da precisión reutilizamos por completo a implementación de inferencia tanto do modelo obxectivo como do preliminar.

Ao escalar desde o conxunto de datos ShareGPT utilizado no artigo orixinal a mostras máis grandes, notamos que a arquitectura de cabeza MTP esbozada e implementada no artigo de EAGLE non puido adestrarse para modelos do tamaño de 70B. A diferenza de ShareGPT, que contiña un maior número de secuencias máis curtas, adestramos nun número lixeiramente menor de consignas substancialmente máis longas. Dado que as cabezas de EAGLE orixinais divergían lixeiramente en estrutura dun transformador típico, reintroducimos algunhas capas de normalización RMS que foramos eliminadas. Descubrimos que isto non só permitía converxer o adestramento, senón que tamén aumentaba a precisión das cabezas nuns poucos puntos porcentuais.

Ampliar

As normas de capa non só facilitan o adestramento, senón que reintroducir as normas tamén é matematicamente intuitivo. As cabezas MTP reutilizan as incrustacións e as proxeccións de logits do modelo obxectivo, xa que poden ser de tamaño substancial (uns 2 GB para Llama 70B). Durante o adestramento, estas están conxeladas e a expectativa é que a capa MTP aprenda a incrustar predicións no mesmo espazo vectorial que o que aprendeu a capa de proxección do modelo orixinal durante o adestramento. Ao eliminar as normas, espérase que un único MLP aprenda a mesma función que un MLP seguido dunha norma, o que dificulta a coincidencia entre os estados ocultos do borrador e os modelos obxectivo.

Inferencia con decodificación especulativa

No motor de inferencia, para xerar tokens para secuencias de entrada, primeiro deben agruparse en lotes de tamaño razoable e, a continuación, deben asignarse páxinas na caché KV para os seguintes tokens. Os tokens de entrada e a información da páxina KV envólvense nun búfer transmitido a todos os rangos paralelos que executan o modelo. Finalmente, os metadatos cópianse na memoria da GPU e o modelo execútase para producir os logits a partir dos cales se mostras o seguinte token.

A diferenza de certas implementacións que acoplan frouxamente un servidor de inferencia preliminar e obxectivo mediante un envoltorio que orquestra as solicitudes entre eles, os nosos pares de borrador-obxectivo están fortemente acoplados e avanzan na xeración ao unísono. A programación de lotes e a asignación de páxinas KV compártense entre os modelos para todas as formas de decodificación especulativa: isto unifica a lóxica que une un modelo co servidor de inferencia global, xa que todos expoñen a mesma interface.

O tempo de execución da inferencia en Perplexity está estruturado ao redor de FlashInfer, que determina os metadatos que se deben construír para configurar e programar o núcleo de atención. Dadas algunhas secuencias de entrada que forman un lote, para o preenchemento, a decodificación ou a verificación, debe realizarse un traballo do lado da CPU para asignar búferes intermedios e encher certos búferes constantes utilizados na atención. Este traballo súmase ao custo da programación de lotes e da asignación de páxinas KV, que tamén incorren en latencias que se deben ocultar para maximizar a utilización da GPU.

Aínda que paralelizamos por completo o traballo do lado da CPU e da GPU para a inferencia sen especulación, atopamos que o equilibrio CPU-GPU para a decodificación especulativa é máis intrincado. O principal desafío xorde do feito de que o número de tokens aceptados determina a lonxitude da secuencia para unha execución posterior, introducindo un punto de sincronización de GPU a CPU difícil de evitar. Experimentamos con diferentes esquemas de programación para ocultar mellor a latencia do traballo da CPU.

Programación de borrador-obxectivo

A pesar de ser máis pequeno que un modelo obxectivo, cando se utiliza un LLM completo como borrador, aínda introduce unha latencia considerable na GPU, proporcionando certa marxe para ocultar operacións custosas de CPU. Dado que os modelos máis pequenos non se benefician do paralelismo de tensores, hai unha desasociación entre o número de rangos nos que se divide un obxectivo e un borrador. Na nosa implementación, o modelo preliminar execútase só no rango líder dun grupo TP.

Ampliar

Como se indicou anteriormente, un paso de decodificación transporta logits á seguinte execución. Isto permítennos solapar unha execución do modelo preliminar co traballo de programación de lotes do lado da CPU. Despois de xuntar o lote, as chamadas repetidas ao mostreador e ao borrador producen os tokens preliminares. En paralelo, o lote de verificación xúntase para o modelo obxectivo e sincronízase cos traballadores paralelos. Os logits obxectivo verifícanse e móstranse para determinar as lonxitudes de secuencia aceptadas. Neste punto, a sincronización de GPU a CPU é necesaria para determinar as lonxitudes de secuencia posteriores. Dado que o modelo preliminar só se executa no nodo líder, o seu lote configúrase secuencialmente e a súa execución iníciase para encher as súas entradas de caché KV co token adicional que produciu o obxectivo. Os logits producidos por esta execución preliminar na execución actual utilizaranse para mostraxe o primeiro token preliminar na execución posterior. O máis importante é que, mentres se executa o borrador, pódese programar o seguinte lote.

Programación MTP para un único token

Aínda que o tempo de execución aínda non ofrece exploración de árbores preliminares ao estilo Eagle, implementamos un caso especial deste esquema, considerando unha secuencia lineal de tokens preliminares producidos por un modelo do tamaño dunha única capa de decodificador de transformador. Este esquema pódese usar para a predición preliminar utilizando os pesos de código aberto de DeepSeek R1. O subcaso de predicir un único token é interesante, xa que as grandes capas MTP logran taxas de aceptación suficientemente altas como para xustificar a súa sobrecarga.

A programación MTP é algo máis complexa, xa que o modelo preliminar é moito máis rápido, ocultando menos latencia do lado da CPU. Ademais, o borrador divídesse xunto co modelo obxectivo, o que require transferencias de memoria compartida para a información por lotes. Unha execución comeza transferindo información por lotes e mostraxe o primeiro token a partir dos logits transportados, de xeito similar ao esquema anterior. A continuación, o obxectivo execútase para validar os tokens, procesando 2 * D tokens, onde D é o tamaño do lote de decodificación. Isto é ideal para microlotes en modelos de mestura de expertos (MoE) a través de interconexións máis lentas como InfiniBand, xa que o lote divídese uniformemente en dúas metades. Os estados ocultos do obxectivo transpórtanse á seguinte execución preliminar, mentres que os logits pásanse ao mostreador para a súa verificación.

Ampliar

Ao realizar unha cantidade limitada de traballo adicional na GPU, evitamos a sincronización de CPU a GPU despois da aceptación da secuencia preliminar. Despois de que se despracen os tokens de entrada dos obxectivos, un núcleo insire os seguintes tokens obxectivo nos seus lugares correspondentes. O borrador vólvese executar coa mesma información por lotes que o obxectivo, enchendo as entradas de caché KV e construíndo os logits e estados ocultos para a seguinte execución, realizando algún traballo redundante en tokens que non foron aceptados. Nestas situacións, a latencia do traballo non utilizado é apenas mensurable debido ao pequeno tamaño do modelo preliminar. En paralelo coa execución preliminar, as lonxitudes das secuencias determínanse na CPU e iníciase a programación do seguinte lote, sen ter que esperar a que remate o traballo da GPU.

A sobrecarga do traballo adicional na capa preliminar non se nota na atención, con todo as capas MLP son máis problemáticas. Dado que as instrucións de multiplicación de matrices axústanse a un límite de 64 ao longo da dimensión do número de tokens, se a duplicación non require significativamente máis bloques, a sobrecarga escóndese. Para secuencias preliminares máis longas, a sobrecarga é máis custosa e o esquema utilizado para os modelos habituales de borrador-obxectivo funciona mellor.

Referencias

Inferencia rápida de transformadores mediante decodificación especulativa

EAGLE: a mostraxe especulativa require repensar a incerteza das características

EAGLE-2: inferencia máis rápida de modelos de linguaxe con árbores preliminares dinámicas

EAGLE-3: escalado da aceleración da inferencia de grandes modelos de linguaxe mediante probas en tempo de adestramento

Medusa: marco de aceleración de inferencia LLM simple con múltiples cabezas de decodificación

Normalización de capas

FlashInfer: motor de atención eficiente e personalizable para a servidume de inferencia LLM

Informe técnico de DeepSeek-V3

Cita

@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}
}