Speculatieve decodering versnelt de generatiesnelheid van grote taalmodellen (LLM's) door een snel en klein ontwerpmodel te gebruiken om voltooiingskandidaten te produceren die worden geverifieerd door het grotere doelmodel.

Onder dit schema, in plaats van dat een dure doelstelling een enkele token produceert, worden er meerdere tokens in één stap uitgezonden. Hier presenteren we de implementatiedetails van verschillende soorten speculatieve decodering, toegepast bij Perplexity om de inter-token latentie op Sonar-modellen te verminderen.

Speculatieve Decodering

Speculatieve Decodering maakt gebruik van de structuur van natuurlijke talen en de autoregressieve aard van transformers om de generatie van tokens te versnellen. Hoewel grotere modellen, zoals Llama-70B, meer kennis bevatten dan kleinere modellen, zoals Llama-1B, presteren ze op sommige eenvoudige taken vergelijkbaar. Deze overlap suggereert dat bepaalde sequenties beter worden gegenereerd door de minder dure modellen, waarbij complexe problemen worden overgelaten aan de grotere modellen. De uitdaging ligt in het bepalen welke voltooiingen beter zijn en of de generatie van het kleinere model van dezelfde kwaliteit is als die van het grotere model.

Gelukkig zijn LLM's autoregressieve transformers: wanneer ze een reeks tokens krijgen, geven ze de kansverdeling van de volgende token. Bovendien geven de logits die zijn afgeleid van de tussenliggende functies die zijn gekoppeld aan de tokens in de invoervolgorde ook aan hoe waarschijnlijk het is dat het model precies die tokens afgeeft. Deze eigenschap maakt speculatie mogelijk: als een reeks tokens wordt gegenereerd door een kleinere en begint bij een invoerprefix, kan het door de grotere worden uitgevoerd om te bepalen hoe goed het aansluit bij het doelmodel. Elk prefix van de kandidaten wordt beoordeeld met een waarschijnlijkheid en de langste boven een acceptatiedrempel wordt gekozen. Als bonus biedt het doelmodel ook gratis een volgende token: als een ontwerpmodel n tokens genereert, kunnen tot n + 1 in één stap worden uitgezonden.

Bij inferentie kan het speculatieve bemonsteringsproces grofweg in 4 fasen worden verdeeld:

1. Vullen: zowel de doel- als de ontwerpmodellen moeten worden uitgevoerd op de invoervolgorde om de KV-cache-invoer te bevolken. Terwijl sommige schema's, zoals Medusa, eenvoudiger dichtheidslagen gebruiken voor voorspelling, richten we ons in deze post op transformer-gebaseerde ontwerpen die hun eigen KV-caches nodig hebben.

2. Ontwerpgeneratie: het ontwerpmodel herhaalt om een aantal vaste tokens te produceren. De ontwerpvolgorde kan lineair zijn of het model kan een boomachtige structuur verkennen tot een bepaalde diepte (EAGLE, Medusa). Hier richten we ons op lineaire sequenties.

3. Acceptatie: het doelmodel werkt de ontwerprange af en bouwt logits overeenkomend met elke ontwerptoken. De lengte van de langste acceptabele sequentie wordt bepaald.

4. Doelgeneratie: aangezien het doel de gegenereerde logits waren, komt bij de niet-overeenkomende positie of het einde van de sequentie de logits overeen met een volgende token. Deze logits kunnen worden bemonsterd om een robuuste token te geven van het doel, en de sequentie af te ronden.

Er bestaan verschillende methoden om speculatieve decodering te implementeren. In deze post zullen we ons richten op de schema's die we hebben gebruikt om de Sonar-modellen te versnellen met behulp van een intern 1B-model, evenals de voorspellingsmechanismen die we aan het ontwikkelen zijn om modellen op de schaal van DeepSeek te versnellen.

Doel-Ontwerp

Speculatieve decodering kan worden bereikt door een bestaand klein LLM als ontwerpmodel te koppelen aan een doelmodel om kandidaatvolgordes te genereren. In de productie hebben we Sonar versneld met een Llama-1B-model dat is fijn afgesteld op hetzelfde dataset als het doel. Hoewel deze aanpak niet vereiste om een ontwerp vanaf nul te trainen, profiteert het kleine model nog steeds aanzienlijk van de KV-cachecapaciteit en introduceert het een lichte overhead bij het prefillen, waardoor TTFT toeneemt.

Onder dit schema speculeert de decoder alleen op decode-only batches, waarbij tokens worden gegenereerd via standaard bemonstering tijdens prefill of op gemengde prefill-decode batches. In de prefill-fase worden de doel-logits onmiddellijk bemonsterd om ook de nieuw gegenereerde token in de KV-cache van het ontwerpmodel te vullen. Het ontwerp wordt nog niet bemonsterd, maar de gegenereerde logits worden meegenomen naar de decode-fase.

In decode, wordt het ontwerpmodel vooruit geschoven, waarbij de top token bij elke fase wordt bemonsterd. Nadat de gewenste ontwerplengte is bereikt, worden de tokens door het doelmodel uitgevoerd om logits te produceren op basis waarvan de sampler de geaccepteerde sequentielengte bepaalt. Acceptatie wordt bepaald door de volledige kansverdelingen van de ontwerpen en het doel te vergelijken. Aangezien het doel altijd een set logits uitvoert die volgt op de geaccepteerde ontwerprange, wordt die bemonsterd om een extra uitvoer te produceren. Aangezien het ontwerpmodel die geaccepteerde token nog niet heeft gezien, wordt het opnieuw uitgevoerd om zijn overeenkomstige KV-cache-invoer voor te bereiden op de volgende decode-stap, en de logits opnieuw mee te nemen.

EAGLE

EAGLE is een speculatieve decodering schema dat meerdere ontwerpranges verkent, gegenereerd via een boomachtige overgang van waarschijnlijke ontwerptokens. Een vaste (EAGLE) of dynamisch gevormde (EAGLE-2) boom wordt verkend met opeenvolgende uitvoeringen van de ontwerptokens, waarbij de Top-K kandidaten op elk knooppunt worden overwogen in plaats van het hoogste scorende token in een lineaire reeks te volgen. De sequenties worden vervolgens beoordeeld en de langste geschikte wordt geselecteerd om verder te gaan, ook een extra token van het doel toe te voegen.

Om een nauwkeurigere voorspelling te bereiken, voorspelt een EAGLE ontwerpmodel niet alleen op basis van tokens, maar ook met behulp van de doelfuncties (laatste laag verstopte staten) van het doelmodel. Het nadeel van EAGLE is de noodzaak om aangepaste, kleine ontwerpmodellen te trainen die nauwkeurig genoeg zijn om geschikte kandidaten binnen een laag latentie-budget te genereren. Meestal is een ontwerpmodel een enkele transformerlaag identiek aan een decoderlaag van het oorspronkelijke model, die nauw gekoppeld is aan het doel door te binden aan zijn embeddings en lm_head projecties. Omdat dit minder KV-cachecapaciteit vereist, heeft EAGLE een kleinere geheugendruks.

Om boomachtige sequenties in het doelmodel te verifiëren, moeten aangepaste attentiemaskers worden gebruikt. Helaas vertraagt het gebruik van een aangepast attentiemasker voor een hele sequentie aanzienlijk de aandacht voor realistische invoerlengtes (tot 50%), waarmee een deel van de versnellingsvoordeel van speculatie teniet wordt gedaan. Om deze reden hebben we nog geen volledige boomverkenning in productie ingezet, maar richten we ons in plaats daarvan op het speciale geval van enkeletokenvoorspelling via MTP-achtige schema's gepresenteerd in het DeepSeek-V3 Technisch Rapport.

MTP

Dit schema is vergelijkbaar met ontwerpmodel-doelmodel decodering, met uitzondering dat verstopte staten naast tokens worden gebruikt voor voorspelling. Er moet iets meer werk worden gedaan in zowel de prefill als decode-fase in vergelijking met reguliere ontwerpmodel-doelmodel speculatie. Het ontwerpmodel gebruikt zowel tokens als verstopte staten: token t_{i+1} wordt bemonsterd vanaf de logits L_i overeenkomend met token t_i, die op hun beurt worden afgeleid van de verstopte staten H_i. Daardoor moeten de invoer token buffers één stap naar links worden verschoven ten opzichte van de verstopte staat vectoren afgegeven door het doel. De onderstaande figuur markeert de correspondenties gebruikt voor training, evenals de verschuiving tijdens inferentie.

De decodeerstroom is vergelijkbaar met ontwerpmodel-doelmodel decodering, met uitzondering dat zowel verstopte staten als logits worden meegenomen. Onze implementatie deelt alle bijbehorende bemonstering en logitverwerkende logits, waarbij alleen de modelvooruitroepen worden gespecialiseerd. Wanneer meerdere tokens worden voorspeld, gebruikt het ontwerpmodel ontwerpverstopte staten voor voorspelling, waarbij ook KV-cache-invoer worden gevuld op basis van zijn eigen functies. Op de lange termijn kan dit de nauwkeurigheid verminderen. Vervolgens, bij het uitvoeren van het ontwerpmodel om de KV-cache-invoer voor de doelvoorspelling te vullen, voeren we het op de hele sequentie uit en nemen de nauwkeurigere doelfuncties als invoer. Omdat deze ontwerpmodellen klein zijn, is de toegevoegde kost van het verwerken van de extra tokens te verwaarlozen.

MTP Hoofdtraining

Om van MTP te kunnen profiteren, hebben we de infrastructuur gebouwd die nodig is om MTP hoofden te trainen die zijn aangesloten op onze fijn afgestelde modellen op Perplexity's datasets, uitgevoerd op één knooppunt met 8xH100 apparaten. Binnen ongeveer een dag kunnen we hoofden bouwen voor modellen variërend van Llama-1B tot Llama-70B en DeepSeek V2-Lite. Voor grotere modellen vertrouwen we op MTP hoofden gebouwd tijdens het fijnstemmingsproces.

Het doel van MTP training is om de ontwerpverstopte staten en logits die zijn geëxtrapoleerd uit de doelverstopte staten overeen te laten komen met de volgende token logits en verstopte staten van het doel. Aangezien inferentie voor verstopte staten duur is, berekenen we ze vooraf met onze inferentie-geoptimaliseerde implementatie van het doelmodel, om tijdens de training te gebruiken. Echter, om de inferentie MTP implementatie te valideren en ervoor te zorgen dat numerieke verschillen door kwantisering of optimalisaties de resultaten niet belemmeren, hergebruiken we voor validatieverlies en nauwkeurigheidschatting volledig de inferentie-implementatie van zowel het doel- als het ontwerpmodel.

Bij het opschalen van het ShareGPT dataset zoals in het oorspronkelijke artikel naar grotere samples, merkten we dat de MTP hoofdarchitectuur die in het EAGLE-artikel werd beschreven en geïmplementeerd, niet kon trainen voor 70B-grootte modellen. In tegenstelling tot ShareGPT dat een groter aantal kortere sequenties bevatte, trainen we op een iets kleiner aantal aanzienlijk langere prompts. Omdat de oorspronkelijke EAGLE hoofden iets afwijkend waren in structuur van een typisch transformer, hebben we enkele RMS-normeringslagen die waren gestript opnieuw geïntroduceerd. We vonden dat dit niet alleen training mogelijk maakte om te convergeren, maar ook de nauwkeurigheid van de hoofden met enkele percentagepunten verhoogde.

Niet alleen vergemakkelijken laag-normen het trainen, maar het opnieuw introduceren van de normen is ook mathematisch intuïtief. MTP hoofden hergebruiken de embeddings en de logit projecties van het doelmodel, omdat ze aanzienlijk in omvang kunnen zijn (ongeveer 2 GB voor Llama 70B). Tijdens de training zijn deze bevroren en wordt verwacht dat de MTP-laag leert om voorspellingen te embedden in dezelfde vectorruimte als wat de projectielaag van het oorspronkelijke model leerde tijdens training. Door de normen weg te laten, wordt van een single MLP verwacht dat deze dezelfde functie leert als een MLP gevolgd door een norm, wat de overeenkomst tussen de verstopte staten van het ontwerp- en het doelmodel belemmert.

Inferentie met Speculatieve Decodering

In de inferentie-engine, om tokens voor invoervolgordes te genereren, moeten ze eerst worden gegroepeerd in redelijke batches en vervolgens moeten pagina's worden toegewezen in de KV-cache voor de volgende tokens. De invoertokens en de KV-pagina-informatie worden vervolgens verpakt in een buffer uitegevoerd naar alle parallelle rangen die het model uitvoeren. Tot slot wordt de metadata gekopieerd naar het GPU-geheugen en wordt het model uitgevoerd om de logits te produceren waaruit de volgende token wordt bemonsterd.

In tegenstelling tot bepaalde implementaties die een ontwerp en doel inferentie-server losjes koppelen via een wrapper die aanvragen ertussen orkestreert, zijn onze ontwerp-doel paren strak gekoppeld en doorlopen ze in harmonie de generatie. Batchplanning en KV pagina toewijzing wordt gedeeld tussen de modellen voor alle vormen van speculatieve decodering: dit verenigt de logica die een model koppelt aan de overkoepelende inferentie-server, aangezien ze allemaal dezelfde interface blootstellen.

De inferentie-runtime bij Perplexity is gevormd rond FlashInfer, dat de metadata bepaalt die moet worden gebouwd om de attentiekernel in te stellen en te plannen. Gezien enkele invoersequenties die een batch vormen, voor vullen, decoderen of verificatie, moet er CPU-werk plaatsvinden om tussenbuffers toe te wijzen en bepaalde constante buffers die worden gebruikt in aandacht te vullen. Dit werk is naast de kost van batchplanning en KV-pagina toewijzing, die ook latenties veroorzaken die verborgen moeten worden om GPU-gebruik maximaal te benutten.

Hoewel we het CPU- en GPU-werk volledig parallel hebben gemaakt voor inferentie zonder speculatie, vonden we dat de CPU-GPU balans voor speculatieve decodering meer ingewikkeld is. De belangrijkste uitdaging komt voort uit het feit dat het aantal geaccepteerde tokens de sequentielengte voor een volgende uitvoering bepaalt, waardoor een moeilijk te vermijden GPU-naar-CPU synchronisatiepunt ontstaat. We hebben geëxperimenteerd met verschillende schema's om de latentie van CPU-werk het beste te verbergen.

Ontwerp-Doel Planning

Ondanks dat een doelmodel kleiner is, levert het hele LLM als ontwerp nog steeds aanzienlijke latentie op de GPU, waardoor enige ruimte vrijkomt om dure CPU-bewerkingen te verbergen. Aangezien kleinere modellen niet profiteren van tensor parallelisme, is er een mismatch tussen het aantal rangen dat een doel en ontwerp over is verspreid. In onze implementatie draait het ontwerpmodel alleen op de leidende rang van een TP-groep.

Zoals eerder aangegeven, draagt een decodeerstap de logits over naar de volgende run. Dit stelt ons in staat om één uitvoering van het ontwerpmodel samen te laten vallen met CPU-side batch planning werk. Nadat de batch is samengesteld, worden herhaalde oproepen naar de sampler en het ontwerp uitgevoerd om de ontwerptokens te produceren. Gelijktijdig wordt de batch voor verificatie samengesteld voor het doelmodel en gesynchroniseerd met de parallelle werkers. De doel-logits worden geverifieerd en bemonsterd om de geaccepteerde sequentielengtes te bepalen. Op dit punt is GPU-naar-CPU synchronisatie nodig om de volgende sequentielengtes te bepalen. Aangezien het ontwerpmodel alleen op de leidende node wordt uitgevoerd, wordt de batch ervan sequentieel samengesteld en de uitvoering ervan gestart om de KV-cache-invoer te vullen met de extra token die het doel heeft geproduceerd. De logits die door deze ontwerpuitvoering worden geproduceerd in de huidige run zullen worden gebruikt om de eerste ontwerptoken in de volgende run te bemonsteren. Het belangrijkste is dat, terwijl het ontwerp wordt uitgevoerd, de volgende batch kan worden gepland.

MTP Planning voor een Enkele Token

Hoewel de runtime nog geen Eagle-stijl ontwerpbomverkenning biedt, hebben we een speciaal geval van dit schema geïmplementeerd, rekening houdend met een lineaire reeks van ontwerptokens geproduceerd door een model ter grootte van een enkele transformer decoderlaag. Dit schema kan worden gebruikt voor ontwerppredictie met behulp van de open-source gewichten van DeepSeek R1. De sub-case van het voorspellen van een enkele token is interessant, aangezien grote MTP-lagen voldoende hoge acceptatiepercentages bereiken om hun overhead te rechtvaardigen.

MTP planning is enigszins ingewikkelder, aangezien het ontwerpmodel veel sneller is, waardoor minder CPU-side latentie verborgen kan worden. Bovendien is het ontwerp samen met het doelmodel verdeeld, waardoor gedeelde geheugenoverdrachten voor batchinformatie vereist zijn. Een run start door batchinformatie over te dragen en de eerste token te bemonsteren vanaf carry-over logits, net als bij het vorige schema. Vervolgens wordt het doel uitgevoerd om tokens te valideren, waarbij 2 * D tokens worden verwerkt, waarbij D de decode batch size is. Dit is ideaal voor micro-batching in Mixture-of-Experts (MoE) modellen over langzamere interconnecties zoals InfiniBand, omdat de batch gelijkmatig in twee helften is gesplitst. De verstopte staten van het doel worden overgedragen naar de volgende ontwerprun, terwijl de logits worden doorgegeven aan de sampler voor verificatie.

Door een beperkte hoeveelheid extra werk op de GPU te doen, vermijden we CPU-GPU synchronisatie na acceptatie van de ontwerpreeks. Nadat de invoertokens van de doelen zijn verschoven, plugt een kernel de volgende doeltokens in hun overeenkomstige locaties. Vervolgens wordt het ontwerp opnieuw uitgevoerd met dezelfde batchinformatie als het doel, waarbij KV-cache-invoer wordt gevuld en de logits en verstopte staten worden opgebouwd voor de volgende run, wat wat overtollig werk doet op tokens die niet zijn geaccepteerd. In deze situaties is de latentie van het ongebruikte werk nauwelijks meetbaar vanwege de kleine omvang van het ontwerpmodel. Gelijktijdig met de ontwerpuitvoering worden sequentielengtes op de CPU bepaald en wordt de planning van de volgende batch gestart, zonder te hoeven wachten op beëindiging van GPU-werk.

De overhead van extra werk in de ontwerplaag is niet merkbaar in aandacht, echter MLP-lagen zijn problematischer. Aangezien matrixvermenigvuldigingsinstructies opvullen naar een grens van 64 langs de dimensie van het aantal tokens, als verdubbeling geen aanzienlijk grotere blokken vereist, wordt de overhead verborgen. Voor langere ontwerprangen is de overhead duurder en werkt het schema dat wordt gebruikt voor reguliere ontwerpmodel-doelmodellen beter.

Referenties

• Snelle Inferentie van Transformers via Speculatieve Decodering

• EAGLE: Speculatieve Bemonstering Vereist Herziening van Kenmerk-Onzekerheid

• EAGLE-2: Snellere Inference van Taalmodellen met Dynamische Ontwerp Bomen

• EAGLE-3: Opschaling van Inference Versnelling van Grote Taalmodellen via Training-Tijd Test

• Medusa: Eenvoudig LLM Inference Versnelling Framework met Meerdere Decodeer Hoofden

• Laagnormalisatie

• FlashInfer: Efficiënte en Aanpasbare Attentie-Engine voor LLM Inference Servering

• DeepSeek-V3 Technisch Rapport