Accelerera Sonar genom spekulation

Spekulativ avkodning ökar genereringshastigheten för stora språkmodeller (LLM:er) genom att använda en snabb och liten utkastmodell för att producera slutförandekandidater som verifieras av den större målmodellen. Enligt detta upplägg, i stället för att en körning av den kostsamma modellen producerar en enda token, genereras flera i ett enda steg. Här presenterar vi implementeringsdetaljerna för olika typer av spekulativ avkodning, tillämpade på Perplexity för att minska latensen mellan tokens på Sonar-modeller.

Spekulativ avkodning

Spekulativ avkodning utnyttjar strukturen hos naturliga språk och den autoregressiva naturen hos transformers för att påskynda tokengenerering. Även om större modeller, såsom Llama-70B, bär på mer kunskap än mindre, såsom Llama-1B, presterar de likartat på vissa enklare uppgifter. Denna överlappning antyder att vissa sekvenser genereras bättre av de billigare modellerna, vilket lämnar komplexa problem till de större. Utmaningen ligger i att avgöra vilka slutföranden som är bättre och om genereringen från den mindre modellen håller samma kvalitet som den från den större.

Lyckligtvis är LLM:er autoregressiva transformers: när de ges en sekvens av tokens matar de ut sannolikhetsfördelningen för nästa token. Dessutom indikerar de logits som härleds från de mellanliggande funktioner som är associerade med tokens i inmatningssekvensen också hur sannolikt det är att modellen ger ut just dessa tokens. Denna egenskap möjliggör spekulation: om en sekvens av tokens genereras av en mindre modell med utgångspunkt från ett inmatningsprefix, kan den köras genom den större för att fastställa hur väl den stämmer överens med målmodellen. Varje prefix för kandidaterna poängsätts med en sannolikhet och det längsta prefixet över ett accepttoxiskt tröskelvärde väljs. Som en bonus tillhandahåller målmodellen dessutom en efterföljande token gratis: om en utkastmodell genererar n tokens kan upp till n + 1 genereras i ett enda steg.

Förstora

Vid inferenstid kan den spekulativa samplingsprocessen delas upp i ungefär 4 steg:

Prefill: både mål- och utkastmodellerna måste köras på inmatningssekvensen för att fylla KV-cacheposterna. Medan vissa upplägg, som Medusa, använder enklare tätt lagrade lager för förutsägelse, fokuserar vi i detta inlägg på transformerbaserade utkast som behöver sina egna KV-cacher.

Utkastgenerering: utkastmodellen itererar för att producera ett antal fasta tokens. Utkastsekvensen kan vara linjär eller så kan modellen utforska en trädliknande struktur upp till ett givet djup (EAGLE, Medusa). Här fokuserar vi på linjära sekvenser.

Acceptans: målmodellen körs på utkastsekvensen och bygger logits som motsvarar varje utkasttoken. Längden på den längsta godtagbara sekvensen fastställs.

Målgenerering: eftersom målet genererade logits motsvarar logitsen vid den felaktiga positionen eller slutet av sekvensen en efterföljande token. Dessa logits kan samplas för att tillhandahålla en robust token från målet, vilket avslutar sekvensen.

Det finns olika metoder för att implementera spekulativ avkodning. I det här inlägget kommer vi att fokusera på de upplägg vi använde för att snabba upp Sonar-modeller med hjälp av en intern 1B-modell, samt de förutsägelseemekanismer vi bygger ut för att påskynda modeller i skala med DeepSeek.

Mål-utkast

Spekulativ avkodning kan uppnås genom att koppla en befintlig liten LLM som utkastmodell till en målmodell för att generera kandidatsekvenser. I produktion har vi accelererat Sonar med en Llama-1B-modell finjusterad på samma dataset som målet. Även om denna metod inte krävde att ett utkast tränades från grunden, använder den lilla modellen fortfarande betydifikativ KV-cachekapacitet och introducerar en viss prefill-overhead, vilket ökar TTFT.

Enligt detta upplägg spekulerar avkodaren enbart på avkodningsspecifika batcher, och genererar tokens genom standardsampling under prefill eller på blandade prefill-avkodningsbatcher. Under prefill-steget samplas målets logits omedelbart för att även fylla i den nyligen genererade token i utkastets KV-cache. Utkastet samplas inte än, men de logits det producerar förs vidare till avkodningssteget.

Förstora

Vid avkodning drivs utkastmodellen framåt och samplar den översta token vid varje steg. När önskad utkastlängd har nåtts körs tokens genom målmodellen för att producera de logits utifrån vilka samplaren identifierar den accepterade sekvenslängden. Acceptans bestäms genom att jämföra de fullständiga sannolikhetsfördelningarna från utkastet och målet. Eftersom målet alltid matar ut en uppsättning logits efter den accepterade utkastsekvensen, samplas denna för att producera en extra utdata. Eftersom utkastmodellen ännu inte har sett den accepterade token körs den om för att fylla dess motsvarande KV-cacheposter som förberedelse för nästa avkodningssteg, varvid logits förs över igen.

EAGLE

EAGLE är ett spekulativt avkodningsupplägg som utforskar flera utkastsekvenser som genereras genom en trädliknande genomgång av sannolika utkasttokens. Ett fast (EAGLE) eller dynamiskt format (EAGLE-2) träd utforskas med hjälp av på varandra följande körningar av utkasttokens, där man tar hänsyn till de Top-K-kandidater vid varje nod istället för att följa den token med högst poäng i en linjär sekvens. Sekvenserna poängsätts sedan och den längsta lämpliga väljs för att fortsätta, varvid en extra token från målet också läggs till.

Förstora

För att uppnå mer exakt förutsägelse förutsäger en EAGLE-utkastmodell inte bara baserat på tokens, utan använder även målmodellens målegenskaper (dolda tillstånd i det sista lagret). Nackdelen med EAGLE är behovet av att träna anpassade, små utkastmodeller som är tillräckligt exakta för att generera lämpliga kandidater inom en låg latensbudget. Typiskt sett är en utkastmodell ett enda transformerlager som är identiskt med ett avkodarlager i originalmodellen, vilket är tätt kopplat till målet genom att knytas till dess inbäddningar och lm_head-projektioner. Eftersom detta kräver mindre KV-cachekapacitet har EAGLE ett lägre minnesfotavtryck.

För att verifiera trädliknande sekvenser i målmodellen måste anpassade uppmärksamhetsmasker (attention masks) användas. Tyvärr saktar användningen av en anpassad uppmärksamhetsmask för en hel sekvens ner uppmärksamheten avsevärt för realistiska inmatningslängder (med upp till 50 %), vilket upphäver en del av den hastighetsökning som kan uppnås genom spekulation. Vi har av denna anledning ännu inte distribuerat fullständig trädutforskning till produktion, utan fokuserar i stället på specialfallet med förutsägelse av enskilda tokens via MTP-liknande upplägg som presenteras i DeepSeek-V3 Technical Report.

MTP

Detta upplägg liknar utkast-mål-avkodning, med undantag för att dolda tillstånd används tillsammans med tokens för förutsägelse. Något mer arbete måste utföras i både prefill- och avkodningsstegen jämfört med vanlig utkast-mål-spekulation. Utkastmodellen använder både tokens och dolda tillstånd: token t_{i+1} samplas från logits L_i som motsvarar token t_i, vilka i sin tur härleds från de dolda tillstånden H_i. Följaktligen måste inmatningstokenbuffertarna skiftas ett steg till vänster i förhållande till de dolda tillståndsvektorer som matas ut av målet. Figuren nedan markerar de överensstämmelser som används för träning samt skiftet under inferens.

Förstora

Avkodningsflödet liknar ganska mycket utkast-mål-avkodning, med undantag för att både dolda tillstånd och logits förs vidare. Vår implementering delar alla associerade samplings- och logitbearbetnings-logits, och specialiserar enbart modellens framåtriktade anrop. När flera tokens förutsägs använder utkastmodellen dolda utkasttillstånd för förutsägelse och fyller även i KV-cacheposter baserat på egna funktioner. På sikt kan detta försämra noggrannheten. Därefter, när utkastmodellen körs för att fylla KV-cacheposten för målförutsägelsen, kör vi den på hela sekvensen och tar de mer exakta dolda måltillstånden som inmatning. Eftersom dessa utkastmodeller är små är den tillkommande kostnaden för att bearbeta de extra tokens försumbar.

Träning av MTP-huvuden

För att dra nytta av MTP byggde vi den infrastruktur som krävs för att träna MTP-huvuden kopplade till våra finjusterade modeller på Perplexity dataset, körande på en nod med 8xH100-enheter. På ungefär en dag kan vi bygga huvudmodeller för modeller som sträcker sig från Llama-1B till Llama-70B och DeepSeek V2-Lite. För större modeller förlitar vi oss på MTP-huvuden som byggts under finjusteringsprocessen.

Målet med MTP-träning är att matcha de dolda utkasttillstånden och logits som extrapolerats från de dolda måltillstånden till nästa tokens logits och dolda tillstånd för målet. Eftersom inferens för dolda tillstånd är kostsam förberäknar vi dem med vår inferensoptimerade implementering av målmodellen för att användas under träning. För att validera MTP-implementeringen för inferens och säkerställa att numeriska skillnader på grund av kvantisering eller optimeringar inte hindrar resultaten, återanvänder vi dock fullt ut inferensimplementeringen för både mål- och utkastmodellerna för valideringsförlust och uppskattning av noggrannhet.

Vid skalning från ShareGPT-datasetet som användes i originalartikeln till större urval märkte vi att MTP-huvudarkitekturen som beskrevs och implementerades i EAGLE-artikeln misslyckades med att tränas för 70B-modeller. Till skillnad från ShareGPT, som innehöll ett större antal kortare sekvenser, tränar vi på ett något mindre antal väsentligt längre prompter. Eftersom de ursprungliga EAGLE-huvudena skilde sig något i struktur från en typisk transformer, återinförde vi några RMS-normaliseringslager som hade tagits bort. Vi fann att detta inte bara gjorde det möjligt för träningen att konvergera, utan det ökade också huvudenas noggrannhet med några procentenheter.

Förstora

Lagersormaliseringar underlättar inte bara träning, utan att återinföra normerna är också matematiskt intuitivt. MTP-huvuden återanvänder målmodellens inbäddningar och logitprojektioner eftersom de kan vara avsevärda i storlek (cirka 2 GB för Llama 70B). Under träning fryses dessa och förväntningen är att MTP-lagret lär sig att bädda in förutsägelser i samma vektorutrymme som projektionslagret i originalmodellen lärde sig under träning. Genom att ta bort normerna förväntas ett enda MLP lära sig samma funktion som ett MLP följt av en norm, vilket förhindrar matchningen mellan de dolda tillstånden hos utkastet och målmodellerna.

Inferens med spekulativ avkodning

I inferensmotorn måste inmatningssekvenserna, för att generera tokens, först grupperas i rimligt stora batcher, varefter sidor måste allokeras i KV-cachen för nästa tokens. Inmatningstokens och KV-sidinformationen packas sedan in i en buffert som sänds ut till alla parallella ranker som kör modellen. Slutligen kopieras metadata till GPU-minnet och modellen exekveras för att producera de logits från vilka nästa token samplas.

Till skillnad från vissa implementeringar som löst kopplar samman en utkast- och målinferensserver via en wrapper som orkestrerar begäranden mellan dem, är våra utkastmålpar tätt sammankopplade och går igenom genereringen i takt. Batchschemaläggning och KV-sidallokering delas mellan modellerna för alla former av spekulativ avkodning: detta förenar den logik som överbryggar en modell med den övergripande inferensservern, eftersom de alla exponerar samma gränssnitt.

Inferenskörningen på Perplexity är formad kring FlashInfer, som bestämmer de metadata som måste byggas för att konfigurera och schemalägga uppmärksamhetskärnan. Givet några inmatningssekvenser som bildar en batch, måste CPU-sidearbete utföras för prefill, avkodning eller verifiering för att allokera mellliggande buffertar och fylla i vissa konstanta buffertar som används vid uppmärksamhet. Detta arbete kommer utöver kostnaden för batchschemaläggning och KV-sidallokering, vilket också medför latenser som måste döljas för att maximera GPU-utnyttjandet.

Även om vi fullständigt parallelliserade CPU- och GPU-sidans arbete för inferens utan spekulation fann vi att CPU-GPU-balansen för spekulativ avkodning är mer invecklad. Den huvudsakliga utmaningen härstammar från det faktum att antalet accepterade tokens avgör sekvenslängden för en efterföljande körning, vilket introducerar en svårundviklig GPU-till-CPU-synkroniseringspunkt. Vi experimenterade med olika schemaläggningsupplägg för att på bästa sätt dölja latensen för CPU-arbetet.

Utkast-mål-schema

Trots att den är mindre än en målmodell introducerar en hel LLM som används som utkast fortfarande avsevärd latens på GPU:n, vilket ger visst utrymme för att dölja dyra CPU-operationer. Eftersom mindre modeller inte drar nytta av tensor-parallellism finns det en missmatchning mellan antalet ranker som ett mål och ett utkast är shardade över. I vår implementering körs utkastmodellen enbart på ledarranken för en TP-grupp.

Förstora

Som angivits tidigare för ett avkodningssteg över logits till nästa körning. Detta gör att vi kan överlappa en körning av utkastmodellen med CPU-sidans batchschemaläggningsarbete. Efter att batchen har satts samman producerar upprepade anrop till samplaren och utkastet utkasttokens. Parallellt sätts batchen för verifiering ihop för målmodellen och synkroniseras med de parallella arbetarna. Målets logits verifieras och samplas för att bestämma de accepterade sekvenslängderna. Vid denna tidpunkt är GPU-till-CPU-synkronisering nödvändig för att fastställa efterföljande sekvenslängder. Eftersom utkastmodellen endast körs på ledarnoden ställs dess batch in sekventiellt och dess exekvering sätts igång för att fylla dess KV-cacheposter med den extra token som målet producerade. De logits som produceras av denna utkastkörning i den aktuella körningen kommer att användas för att sampla den första utkasttoken i den efterföljande körningen. Viktigast av allt är att nästa batch kan schemaläggas medan utkastet körs.

MTP-schema för en enskild token

Även om körtiden ännu inte tillhandahåller utskatsträdutforskning i Eagle-stil implementerade vi ett specialfall av detta upplägg med beaktande av en linjär sekvens av utkasttokens som producerats av en modell i storleken av ett enda transformeravkodarlager. Detta upplägg kan användas för utkastförutsägelse med hjälp av öppen källkods-vikterna för DeepSeek R1. Subfallet att förutsäga en enda token är intressant eftersom stora MTP-lager uppnår tillräckligt höga acceptansgrader för att motivera sin overhead.

MTP-schemaläggning är något mer komplex eftersom utkastmodellen är mycket snabbare och döljer mindre CPU-sidans latens. Dessutom shardas utkastet tillsammans med målmodellen, vilket kräver delade minnesöverföringar för batchinformation. En körning startar med att överföra batchinformation och sampla den första token från medföljande logits, på samma sätt som i det föregående upplägget. Därefter körs målet för att validera tokens och bearbetar 2 * D tokens, där D är avkodningsbatchstorleken. Detta är idealiskt för mikro-batching i Mixture-of-Experts (MoE)-modeller över långsammare interkonnektioner som InfiniBand, eftersom batchen delas upp jämnt i två hälfter. Målets dolda tillstånd förs över till nästa utkastkörning, medan logits skickas in i samplaren för verifiering.

Förstora

Genom att utföra en begränsad mängd extra arbete på GPU:n undviker vi GPU-till-CPU-synkronisering efter acceptans av utkastsekvensen. Efter att målens inmatningstokens har skiftats in placerar en kärna in nästa måltokens på deras motsvarande platser. Utkastet körs sedan om med samma batchinformation som målet, fyller i KV-cacheposter och bygger logits och dolda tillstånd för nästa körning, varvid visst redundant arbete utförs på tokens som inte accepterades. I dessa situationer är latensen för det oanvända arbetet knappt mätbar på grund av utkastmodellens lilla storlek. Parallellt med utkastkörningen fastställs sekvenslängder på CPU:n och schemaläggningen av nästa batch sätts igång utan att man behöver vänta på att GPU-arbetet ska avslutas.

Overheaden för det extra arbetet i utkastlagret är inte märkbar vid uppmärksamhet, men MLP-lager är mer problematiska. Eftersom matrismultiplikationsinstruktioner fyller ut till en gräns på 64 längs dimensionen för antalet tokens, döljs overheaden om dubbleringen inte kräver avsevärt fler block. För längre utkastsekvenser är overheaden dyrare och upplägget som används för vanliga utkast-mål-modeller fungerar bättre.

Referenser

Snabb inferens från transformers via spekulativ avkodning

EAGLE: Spekulativ sampling kräver omprövning av osäkerhet i funktioner

EAGLE-2: Snabbare inferens av språkmodeller med dynamiska utkastträd

EAGLE-3: Skalning av inferensacceleration för stora språkmodeller via test under träning

Medusa: Enkelt ramverk för LLM-inferensacceleration med flera avkodningshuvuden

Lagersnormalisering

FlashInfer: Effektiv och anpassningsbar uppmärksamhetsmotor för LLM-inferenstjänster

DeepSeek-V3 Technical Report

Citering

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