Acceleració de Sonar mitjançant l'especulació

La descodificació especulativa accelera la velocitat de generació dels models de llenguatge grans (LLM) utilitzant un model preliminar ràpid i petit per produir candidats de compleció que són verificats pel model objectiu, que és més gran. Sota aquest esquema, en lloc d'una execució del model objectiu (que és costosa) per produir un sol token, se n'emeten diversos en un sol pas. Aquí presentem els detalls d'implementació de diversos tipus de descodificació especulativa, aplicats a Perplexity per reduir la latència entre tokens als models Sonar.

Descodificació especulativa

La descodificació especulativa aprofita l'estructura de les llengües naturals i la naturalesa autoregressiva dels transformadors per accelerar la generació de tokens. Tot i que els models més grans, com ara Llama-70B, tenen més coneixement que els més petits, com ara Llama-1B, en algunes tasques més senzilles funcionen de manera similar. Aquesta superposició suggereix que determinades seqüències es generen millor amb els models menys costosos, deixant els problemes complexos als més grans. El repte rau a determinar quines complecions són millors i si la generació del model més petit té la mateixa qualitat que la del més gran.

Afortunadament, els LLM són transformadors autoregressius: quan se'ls dóna una seqüència de tokens, mostren la distribució de probabilitat del token següent. A més, els lògits derivats de les característiques intermèdies associades als tokens de la seqüència d'entrada també indiquen com de probable és que el model emeti aquests tokens exactes. Aquesta propietat permet l'especulació: si un model més petit genera una seqüència de tokens a partir d'un prefix d'entrada, es pot executar a través del model més gran per determinar com s'alinea amb el model objectiu. Cada prefix dels candidats es puntua amb una probabilitat i s'escull el més llarg que supera un llindar d'acceptació. Com a bonificació, el model objectiu també proporciona un token posterior de franc: si un model preliminar genera n tokens, se'n poden emetre fins a n + 1 en un sol pas.

Amplia

En el moment de la inferència, el procés de mostreig especulatiu es pot dividir aproximadament en 4 etapes:

Emplenament previ (prefill): tant el model objectiu com el preliminar s'han d'executar a la seqüència d'entrada per omplir les entrades de la memòria cau KV. Tot i que alguns esquemes, com ara Medusa, utilitzen capes denses més senzilles per a la predicció, en aquesta publicació ens centrem en esborranys basats en transformadors que necessiten les seves pròpies memòries cau KV.

Generació de l'esborrany: el model preliminar itera per produir un nombre de tokens fixos. La seqüència de l'esborrany pot ser lineal o el model pot explorar una estructura en forma d'arbre fins a una profunditat determinada (EAGLE, Medusa). Aquí ens centrem en les seqüències lineals.

Acceptació: el model objectiu s'executa a la seqüència de l'esborrany, construint lògits que corresponen a cada token de l'esborrany. Es determina la longitud de la seqüència acceptable més llarga.

Generació de l'objectiu: com que l'objectiu ha generat lògits, a la posició on no coincideixen o al final de la seqüència, els lògits corresponen a un token posterior. Aquests lògits es poden mostrejar per proporcionar un token robust del model objectiu, completant la seqüència.

Existeixen diversos mètodes per implementar la descodificació especulativa. En aquesta publicació, ens centrem en els esquemes que hem utilitzat per accelerar els models Sonar utilitzant un model 1B intern, així com els mecanismes de predicció que estem construint per accelerar models a l'escala de DeepSeek.

Objectiu-Esborrany

La descodificació especulativa es pot aconseguir acoblant un LLM petit existent com a model preliminar a un model objectiu per generar seqüències candidates. En producció, hem accelerat Sonar utilitzant un model Llama-1B ajustat amb el mateix conjunt de dades que l'objectiu. Tot i que aquest enfocament no requeria entrenar un esborrany des de zero, el model petit encara utilitza una capacitat significativa de la memòria cau KV i introdueix una lleugera sobrecàrrega d'emplenament previ, augmentant el TTFT.

Sota aquest esquema, el descodificador només especula amb lots només de descodificació, generant tokens mitjançant el mostreig estàndard durant l'emplenament previ o en lots mixts d'emplenament previ i descodificació. A l'etapa d'emplenament previ, els lògits de l'objectiu es mostregen immediatament per emplenar també el token acabat de generar a la memòria cau KV de l'esborrany. L'esborrany encara no es mostreja, però els lògits que produeix es transfereixen a l'etapa de descodificació.

Amplia

En la descodificació, el model preliminar avança, mostrejant el token principal a cada etapa. Un cop s'arriba a la longitud de l'esborrany desitjada, els tokens s'executen a través del model objectiu per produir els lògits a partir dels quals el mostrejador identifica la longitud de la seqüència acceptada. L'acceptació es determina comparant les distribucions de probabilitat completes de l'esborrany i l'objectiu. Com que l'objectiu sempre emet un conjunt de lògits després de la seqüència de l'esborrany acceptada, aquest es mostreja per produir una sortida addicional. Com que el model preliminar encara no ha vist aquest token acceptat, s'executa de nou per omplir les entrades corresponents de la memòria cau KV en preparació per al següent pas de descodificació, transferint de nou els lògits.

EAGLE

EAGLE és un esquema de descodificació especulativa que explora múltiples seqüències d'esborrany, generades mitjançant un recorregut en forma d'arbre de tokens d'esborrany probables. S'explora un arbre fix (EAGLE) o de forma dinàmica (EAGLE-2) utilitzant execucions consecutives dels tokens de l'esborrany, tenint en compte els candidats Top-K a cada node en lloc de seguir el token amb la puntuació més alta en una seqüència lineal. A continuació, les seqüències es puntuen i se selecciona la més adequada per continuar, afegint també un token addicional de l'objectiu.

Amplia

Per tal d'aconseguir una predicció més precisa, un model preliminar EAGLE prediu no només basant-se en els tokens, sinó també utilitzant les característiques de l'objectiu (estats ocults de l'última capa) del model objectiu. El desavantatge d'EAGLE és la necessitat d'entrenar models preliminars petits i personalitzats que siguin prou precisos per generar candidats adequats dins d'un pressupost de latència baix. Normalment, un model preliminar és una sola capa de transformador idèntica a una capa de descodificador del model original, que està estretament acoblada a l'objectiu enllaçant-la amb les seves incrustacions i projeccions de lm_head. Com que això requereix menys capacitat de memòria cau KV, EAGLE té una empremta de memòria menor.

Per verificar seqüències en forma d'arbre al model objectiu, s'han d'utilitzar màscares d'atenció personalitzades. Malauradament, utilitzar una màscara d'atenció personalitzada per a tota una seqüència înc slow down (alenteix significativament) l'atenció per a longituds d'entrada realistes (fins a un 50%), anul·lant part de l'acceleració aconseguida mitjançant l'especulació. Per aquest motiu, encara no hem desplegat l'exploració completa d'arbres a producció, sinó que ens centrem en el cas especial de la predicció d'un sol token mitjançant esquemes similars a MTP presentats a l'informe tècnic de DeepSeek-V3.

MTP

Aquest esquema és similar a la descodificació d'esborrany-objectiu, amb l'excepció que s'utilitzen estats ocults juntament amb els tokens per a la predicció. S'ha de fer una mica més de feina tant en les etapes d'emplenament previ com de descodificació en comparació amb l'especulació regular d'esborrany-objectiu. El model preliminar utilitza tant tokens com estats ocults: el token t_{i+1} es mostreja a partir dels lògits L_i corresponents al token t_i, que al seu aldar es deriven dels estats ocults H_i. En conseqüència, els búfers de tokens d'entrada s'han de desplaçar un pas a l'esquerra en relació amb els vectors d'estat ocult sortits per l'objectiu. La figura següent marca les correspondències utilitzades per a l'entrenament, així com el desplaçament durant la inferència.

Amplia

El flux de descodificació és bastant similar a la descodificació d'esborrany-objectiu, amb l'excepció que tant els estats ocults com els lògits es transfereixen. La nostra implementació comparteix tot el mostreig associat i el processament de lògits, especialitzant només les invocacions cap endavant del model. Quan es prediuen diversos tokens, el model preliminar utilitza estats ocults de l'esborrany per a la predicció, omplint també les entrades de la memòria cau KV basant-se en les seves pròpies característiques. A la llarga, això pot degradar la precisió. Posteriorment, quan s'executa el model preliminar per omplir l'entrada de la memòria cau KV per a la predicció de l'objectiu, l'executem en tota la seqüència prenent els estats ocults de l'objectiu, que són més precisos, com a entrades. Com que aquests models preliminars són petits, el cost afegit de processar els tokens addicionals és negligible.

Entrenament de caps MTP

Per tal de beneficiar-nos de MTP, vam construir la infraestructura necessària per entrenar caps MTP adjunts als nostres models ajustats en els conjunts de dades de Perplexity, executant-nos en un node amb dispositius 8xH100. En aproximadament un dia, podem construir caps per a models que van des de Llama-1B fins a Llama-70B i DeepSeek V2-Lite. Per als models més grans, confiem en els caps MTP construïts durant el procés d'ajust.

L'objectiu de l'entrenament MTP és fer coincidir els estats ocults i els lògits de l'esborrany extrapolats dels estats ocults de l'objectiu amb els lògits del token següent i els estats ocults de l'objectiu. Com que la inferència per als estats ocults és costosa, els precalculam utilitzant la nostra implementació optimitzada per a la inferència del model objectiu, per ser utilitzats durant l'entrenament. No obstant això, per validar la implementació MTP de inferència i assegurar-nos que les diferències numèriques degudes a la quantització o les optimitzacions no obstaculitzin els resultats, per a la pèrdua de validació i l'estimació de precisió reutilitzem completament la implementació de inferència tant del model objectiu com del preliminar.

Quan vam escalar des del conjunt de dades ShareGPT utilitzat a l'article original fins a mostres més grans, vam notar que l'arquitectura del cap MTP descrita i implementada a l'article d'EAGLE no va aconseguir entrenar-se per a models de la mida de 70B. A diferència de ShareGPT, que contenia un nombre més gran de seqüències més curtes, nosaltres entrenem amb un nombre lleugerament menor d'indicacions substancialment més llargues. Com que els caps originals d'EAGLE divergien lleugerament en l'estructura d'un transformador típic, vam reintroduir algunes capes de normalització RMS que s'havien eliminat. Vam descobrir que això no només permetia que l'entrenament convergís, sinó que també augmentava la precisió dels caps uns quants punts percentuals.

Amplia

Les normes de capa no només faciliten l'entrenament, sinó que reintroduir les normes també és intuïtiu des del punt de vista matemàtic. Els caps MTP reutilitzen les incrustacions i les projeccions de lògits del model objectiu, ja que poden ser de mida substancial (uns 2 GB per a Llama 70B). Durant l'entrenament, aquestes es congelen i s'espera que la capa MTP aprengui a incrustar prediccions al mateix espai vectorial que el que la capa de projecció del model original va aprendre durant l'entrenament. En eliminar les normes, s'espera que un sol MLP aprengui la mateixa funció que un MLP seguit d'una norma, la qual cosa dificulta la coincidència entre els estats ocults de l'esborrany i dels models objectiu.

Inferència amb descodificació especulativa

Al motor de inferència, per generar tokens per a seqüències d'entrada, primer s'han d'agrupar en lots de mida raonable, i després s'han d'assignar pàgines a la memòria cau KV per als tokens següents. A continuació, els tokens d'entrada i la informació de la pàgina KV s'empaqueten en un búfer que es difon a tots els rangs paral·lels que executen el model. Finalment, les metadades es copien a la memòria de la GPU i s'executa el model per produir els lògits a partir dels quals es mostreja el token següent.

A diferència de determinades implementacions que acoblen de manera fluixa un servidor de inferència preliminar i objectiu mitjançant un connector intermedi (wrapper) que orquestra les sol·licituds entre ells, les nostres parelles d'esborrany i objectiu estan estretament acoblades i passen per la generació a l'uníson. La planificació de lots i l'assignació de pàgines KV es comparteixen entre els models per a totes les formes de descodificació especulativa: això unifica la lògica que connecta un model amb el servidor de inferència global, ja que tots exposeu la mateixa interfície.

El temps d'execució de la inferència a Perplexity s'estructura al voltant de FlashInfer, que determina les metadades que s'han de construir per configurar i planificar el nucli d'atenció. Donades algunes seqüències d'entrada que formen un lot, per a l'emplenament previ, la descodificació o la verificació, cal fer feina al costat de la CPU per assignar búfers intermedis i emplenar determinats búfers constants utilitzats en l'atenció. Aquest treball s'afegeix al cost de la planificació de lots i l'assignació de pàgines KV, que també incorren en latències que s'han d'ocultar per maximitzar la utilització de la GPU.

Tot i que vam paral·lelitzar completament el treball de la CPU i la GPU per a la inferència sense especulació, vam descobrir que l'equilibri entre la CPU i la GPU per a la descodificació especulativa és més complex. El principal repte sorgeix del fet que el nombre de tokens acceptats determina la longitud de la seqüència per a una execució posterior, la qual cosa introdueix un punt de sincronització entre la GPU i la CPU que és difícil d'evitar. Vam experimentar amb diferents esquemes de planificació per tal d'amagar millor la latència del treball de la CPU.

Planificació d'esborrany-objectiu

Tot i ser més petit que un model objectiu, quan s'utilitza un LLM sencer com a esborrany, encara introdueix una latència considerable a la GPU, proporcionant un cert marge per amagar operacions costoses de la CPU. Com que els models més petits no es beneficien del paralel·lisme de tensors, hi ha una manca de coincidència entre el nombre de rangs en què es divideix un objectiu i un esborrany. En la nostra implementació, el model preliminar s'executa només al rang principal d'un grup TP.

Amplia

Tal com s'ha indicat abans, un pas de descodificació transfereix els lògits a la següent execució. Això ens permet superposar una execució del model preliminar amb el treball de planificació de lots del costat de la CPU. Un cop reunit el lot, les trucades repetides al mostrejador i a l'esborrany produeixen els tokens de l'esborrany. En paral·lel, el lot de verificació es reuneix per al model objectiu i se sincronitza amb els treballadors en paral·lel. Els lògits de l'objectiu es verifiquen i es mostregen per determinar les longituds de seqüència acceptades. En aquest punt, la sincronització de GPU a CPU és necessària per determinar les longituds de seqüència posteriors. Com que el model preliminar només s'executa al node principal, el seu lot es configura seqüencialment i s'inicia la seva execució per omplir les entrades de la memòria cau KV amb el token addicional que ha produït l'objectiu. Els lògits produïts per aquesta execució de l'esborrany a l'execució actual s'utilitzaran per mostrejar el primer token de l'esborrany a l'execució posterior. El més important és que, mentre s'executa l'esborrany, es pot planificar el lot següent.

Planificació MTP per a un sol token

Tot i que el temps d'execució encara no ofereix l'exploració d'arbres d'esborrany a l'estil d'Eagle, hem implementat un cas especial d'aquest esquema, considerant una seqüència lineal de tokens d'esborrany produïts per un model de la mida d'una sola capa de descodificador de transformador. Aquest esquema es pot utilitzar per a la predicció d'esborranys utilitzant els pesos de codi obert de DeepSeek R1. El subcas de predir un sol token és interessant, ja que les capes MTP grans assoleixen taxes d'acceptació prou altes per justificar la seva sobrecàrrega.

La planificació MTP és una mica més complexa, ja que el model preliminar és molt més ràpid, amagant menys latència del costat de la CPU. A més, l'esborrany es divideix juntament amb el model objectiu, la qual cosa requereix transferències de memòria compartida per a la informació dels lots. Una execució comença transferint la informació del lot i mostrejant el primer token a partir dels lògits transferits, de manera similar a l'esquema anterior. A continuació, s'executa l'objectiu per validar els tokens, processant 2 * D tokens, on D és la mida del lot de descodificació. Això és ideal per a lots petits (micro-batching) en models de mescla d'experts (MoE) a través d'interconnexions més lentes com InfiniBand, ja que el lot es divideix uniformement en dues meitats. Els estats ocults de l'objectiu es transfereixen a la següent execució de l'esborrany, mentre que els lògits es passen al mostrejador per a la verificació.

Amplia

Realitzant una quantitat limitada de treball addicional a la GPU, evitem la sincronització de CPU a GPU després de l'acceptació de la seqüència de l'esborrany. Després de desplaçar els tokens d'entrada dels objectius, un nucli connecta els següents tokens de l'objectiu a les sevesubicacions corresponents. A continuació, l'esborrany s'executa de nou amb la mateixa informació de lot que l'objectiu, omplint les entrades de la memòria cau KV i construint els lògits i els estats ocults per a la següent execució, fent una mica de feina redundant en els tokens que no han estat acceptats. En aquestes situacions, la latència del treball no utilitzat amb prou feines es pot mesurar a causa de la petita mida del model preliminar. En paral·lel amb l'execució de l'esborrany, les longituds de seqüència es determinen a la CPU i s'inicia la planificació del següent lot, sense haver d'esperar que acabi el treball de la GPU.

La sobrecàrrega del treball addicional a la capa de l'esborrany no es nota en l'atenció, no obstant això, les capes MLP són més problemàtiques. Com que les instruccions de multiplicació de matrius s'omple fins a un límit de 64 al llarg de la dimensió del nombre de tokens, si la duplicació no requereix blocs significativament més grans, la sobrecàrrega queda amagada. Per a seqüències d'esborrany més llargues, la sobrecàrrega és més costosa i l'esquema utilitzat per als models regulars d'esborrany-objectiu funciona millor.

Referències

Inferència ràpida a partir de transformadors mitjançant la descodificació especulativa

EAGLE: El mostreig especulatiu requereix repensar la incertesa de les característiques

EAGLE-2: Inferència més ràpida de models de llenguatge amb arbres d'esborrany dinàmics

EAGLE-3: Ampliació de l'acceleració de la inferència de grans models de llenguatge mitjançant proves en temps d'entrenament

Medusa: Marc simple d'acceleració de la inferència d'LLM amb múltiples caps de descodificació

Normalització de capes

FlashInfer: Motor d'atenció eficient i personalitzable per al servei de inferència d'LLM

Informe tècnic de DeepSeek-V3

Citació

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