Speculating Experts Accelerates Inference for Mixture-of-Experts

Dit paper introduceert een expert-prefetching-methode die interne modelrepresentaties gebruikt om toekomstige experts te voorspellen, waardoor CPU-GPU-overdrachten tijdens inferentie van Mixture-of-Experts-modellen overlappen met berekeningen en de tijd per output-token tot 14% wordt verkort.

De kern

Stel je voor dat je een gigantische bibliotheek hebt met miljoenen boeken (de kennis van een kunstmatige intelligentie), maar je kunt er maar een paar tegelijk op je bureau houden. De rest staat in een enorme kelder (het CPU-geheugen).

Elke keer als je een vraag stelt, moet de bibliothecaris (de AI) snel een specifiek boek uit de kelder halen, het op je bureau leggen, het lezen, en dan weer terugbrengen voordat hij de volgende vraag kan beantwoorden. Dit heen-en-weer slepen kost enorm veel tijd. De AI zit vaak te wachten op de bibliothecaris in plaats van te denken. Dit is precies het probleem bij moderne, slimme AI-modellen die "Mixture-of-Experts" (MoE) worden genoemd.

Het probleem: De "Wachtrij"
In deze modellen zijn er duizenden kleine experts (boeken) die elk een specifiek onderwerp kennen. De AI moet voor elk woord dat hij zegt, beslissen welk expert hij nodig heeft. Omdat de computergeheugen (GPU) te klein is voor alle experts, moeten ze worden opgeslagen in de langzamere CPU-geheugen.

• Huidige situatie: De AI denkt na, stopt, wacht tot het juiste boek uit de kelder is gehaald, leest het, en gaat dan pas verder. Dit wachten maakt het hele proces traag.

De oplossing: De "Vooruitziende Blik"
De auteurs van dit paper hebben een slimme truc bedacht: Speculeren (voorspellen).
In plaats van te wachten tot de AI precies weet welk boek hij nodig heeft, kijken ze naar de huidige gedachten van de AI om te raden welk boek hij binnenkort nodig zal hebben.

Hier is hoe het werkt, met een paar creatieve vergelijkingen:

1. De "Quasi-Verborgen Staat" (De Voorspeller)

Stel je voor dat de AI een kok is die een recept maakt.

• Normaal: De kok wacht tot hij de pan heeft leeggemaakt, om dan pas te beslissen of hij peper of zout nodig heeft. Dan rent hij naar de voorraadkast (CPU) om het te halen.
• Met deze nieuwe methode: De kok kijkt naar de geur en de damp van de pan terwijl hij nog aan het koken is. Hij ziet dat de soep dikker wordt en denkt: "Ah, binnen een paar seconden ga ik waarschijnlijk zout nodig hebben."
• De actie: Terwijl de kok nog steeds roert (rekenen), stuurt hij alvast een boodschapper naar de voorraadkast om het zout op te halen. Als de kok het zout echt nodig heeft, ligt het er al klaar. Als hij het niet nodig heeft, is dat ook niet erg; de boodschapper heeft gewoon even gelopen.

In technische termen gebruiken ze een "quasi-verborgen staat". Dit is een soort voorspelling gebaseerd op de huidige input, die zegt: "De volgende stap zal waarschijnlijk Expert A of Expert B nodig hebben."

2. Het Overlappen van Tijden (De Magische Dubbelgang)

Het echte geheim zit in het overlappen.

• Vroeger: Rekenen -> Wachten -> Rekenen -> Wachten.
• Nu: Rekenen (terwijl de boodschapper de volgende boeken haalt) -> Rekenen (terwijl de volgende boeken worden gehaald).

De tijd die normaal wordt verspild aan het wachten op de data, wordt nu gebruikt om de volgende data alvast op te halen. Het is alsof je een trein hebt die niet stopt om passagiers in te laten stappen, maar waar de passagiers al in de trein staan voordat de trein überhaupt stopt.

3. Wat als de voorspelling fout is? (De "Gok")

Soms raadt de AI verkeerd. Hij denkt dat hij zout nodig heeft, maar hij heeft juist suiker nodig.

• Oude methode: Als je een fout raadt, moet je de suiker alsnog uit de kelder halen. Je bent tijd kwijt.
• Nieuwe methode: De auteurs ontdekten iets verrassends: Het maakt vaak niet uit als je een fout raadt.
  • Als de AI "zout" gebruikt in plaats van "suiker", is het resultaat vaak nog steeds goed genoeg om de zin te begrijpen. De AI is zo slim dat hij de fout kan compenseren.
  • Zelfs als de voorspelling niet 100% klopt, is het sneller dan wachten. Het is beter om een beetje onnauwkeurig te zijn en snel te zijn, dan perfect te zijn en traag.

4. De "Slimme Assistent" (Neurale Schatting)

Bij sommige modellen (zoals Qwen) is de voorspelling in de beginfase wat onzekerder. De geur van de pan is nog niet duidelijk genoeg.
Om dit op te lossen, hebben ze een heel klein, slim hulpmiddel (een "neuraal schatter") getraind. Dit is als een assistent die alleen kijkt naar de eerste paar stappen van het recept en zegt: "Zorg dat je zeker de juiste kruiden hebt, want hier maken mensen vaak fouten."
Dit helpt om de voorspellingen in de kritieke beginfase nauwkeuriger te maken, zonder dat het systeem traag wordt.

Het Resultaat

Door deze techniek toe te passen, kunnen deze enorme AI-modellen 5% tot 14% sneller werken op hardware die niet superkrachtig is (zoals een gewone gaming-PC).

• Het betekent dat je een slimme AI op je eigen laptop kunt draaien zonder dat het uren duurt om een antwoord te krijgen.
• Het maakt de "grote" AI-modellen toegankelijk voor gewone mensen, omdat je niet meer een dure server nodig hebt om ze snel te laten draaien.

Samenvattend:
De auteurs hebben een manier gevonden om de AI te leren vooruit te denken. In plaats van te wachten tot hij weet wat hij nodig heeft, haalt hij het alvast op terwijl hij nog aan het werk is. Zelfs als hij een beetje verkeerd raadt, is het resultaat nog steeds goed, en vooral: het is veel sneller. Het is alsof je een auto hebt die niet stopt om te tanken, maar waar de brandstof al in de tank wordt gepompt terwijl je nog rijdt.

Technische samenvatting

Probleemstelling

Mixture-of-Experts (MoE) modellen zijn populair geworden voor het schalen van Large Language Models (LLM's) omdat ze een groot aantal parameters kunnen hebben met een lage rekentijd per token (door slechts een subset van experts te activeren). Echter, in inferentie-scenario's met beperkt geheugen (bijvoorbeeld op consumentenhardware of single-GPU systemen), moeten de expert-weights worden uitgeladen (offloaded) naar het CPU-geheugen.

Dit creëert een zware I/O-bottleneck:

• Tijdens de decode-fase moeten de benodigde expert-weights continu van CPU naar GPU worden overgebracht.
• Deze CPU-GPU-transfers domineren de tijd per output-token (TPOT). Bijvoorbeeld, bij Qwen3-30B-A3B op een A6000 GPU maken transfers 84-88% van de TPOT uit, terwijl de daadwerkelijke berekening slechts 8-13% beslaat.
• Bestaande oplossingen zoals "on-demand loading" (wachten tot de router een expert selecteert en deze dan laden) laten de GPU vaak inactief wachten op data, wat de doorvoer drastisch verlaagt.

Methodologie

De auteurs stellen een expert prefetching-schema voor dat toekomstige expert-selecties voorspelt op basis van interne modelrepresentaties, zodat geheugenoverdrachten kunnen overlappen met berekeningen.

1. Signalen voor Voorspelling:
In plaats van te wachten op de echte router-output van de huidige laag, gebruiken ze interne signalen om de experts van de volgende laag ($l+1$) te voorspellen:

• Default Vector ($d_l$): Een offline berekende vector die de gemiddelde activatie van een expert weergeeft.
• Quasi-Hidden State ($q_l$): Een benadering van de invoer voor de router van de volgende laag. Deze wordt berekend als: $q_l = LN_{l+1}(d_l + r_l)$, waarbij $r_l$ de residual na attention is en $LN$ normalisatie.
• Resultaat: De $q_l$ bevat een "expert-conditioned bias" die helpt bij het voorspellen van de drift tussen de huidige en volgende laag. Voor sommige modellen (zoals GPT-OSS) is dit een sterk signaal; voor andere (zoals Qwen3-30B-A3B) is de drift in vroege lagen groter, wat de voorspellingsnauwkeurigheid verlaagt.

2. Speculatieve Executie:
In tegenstelling tot eerdere werken waar voorspelde experts slechts als "cache hints" dienen (en bij fouten opnieuw geladen moeten worden), voert het systeem de voorspelde experts daadwerkelijk uit.

• Als de voorspelling correct is, wordt de berekening direct uitgevoerd.
• Als de voorspelling fout is, wordt aangenomen dat de impact op de uiteindelijke taaknauwkeurigheid beperkt blijft, omdat de hoogst-gewogen experts (die het meeste bijdragen) vaak wel correct worden voorspeld.
• Dit elimineert de noodzaak om bij een "miss" de echte experts opnieuw te laden, waardoor de overlap tussen CPU-GPU transfers en GPU-berekening maximaal wordt.

3. Implementatie:

• Geïntegreerd in de YALIS inferentie-engine.
• Gebruikt asynchrone transfers (non-blocking CPU→GPU copy) die overlappen met de compute-stream van de huidige laag.
• Implementeert een "wait-and-prefetch" logica: terwijl de GPU de huidige laag berekent, worden de weights voor de volgende laag alvast naar de GPU gestuurd.

4. Verbetering voor hoge drift-lagen (Neurale Estimators):
Voor modellen waar de voorspelling op basis van de quasi-hidden state onnauwkeurig is (vooral in vroege lagen), trainen de auteurs een lichtgewicht neurale estimator.

• Deze estimator leert een mapping van de quasi-hidden state naar de router-logits van de volgende laag.
• Hij wordt getraind via distillatie met weinig tokens (bijv. 4M tokens) en vervangt de router in kritieke lagen om de "hit rate" te verhogen.

Kernbijdragen

1. Parameter-vrije prefetching: Identificatie van interne representaties (quasi-hidden state) die toekomstige routing-beslissingen betrouwbaar voorspellen zonder extra training voor de meeste moderne MoE-architecturen.
2. Speculatieve uitvoering met behoud van nauwkeurigheid: Demonstratie dat het uitvoeren van voorspelde experts (in plaats van ze te negeren bij fouten) de downstream-takenauwkeurigheid behoudt, waardoor de overlap tussen compute en memory transfer optimaal blijft.
3. Geoptimaliseerde implementatie: Integratie in een open-source engine (YALIS) die TPOT met 5-14% reduceert ten opzichte van on-demand loading.
4. Lichtgewicht neurale estimators: Voor architecturen met hoge representatieve drift, wordt een kleine estimator voorgesteld die de voorspellingsnauwkeurigheid aanzienlijk verbetert zonder significante runtime-overhead.

Resultaten

• Snelheidswinst: De methode bereikt een 5-14% reductie in Time Per Output Token (TPOT) ten opzichte van on-demand loading.
  • Op zwakkere GPU's (A6000) is de winst groter (tot 14%) omdat de kopietijd daar dominant is.
  • Op krachtigere GPU's (A100, GH200) is de winst iets lager (5-8%) omdat de berekening sneller is en de kopietijd relatief minder dominant wordt, maar er is nog steeds een duidelijke verbetering.
• Nauwkeurigheid:
  • Voor GPT-OSS modellen blijft de nauwkeurigheid op benchmarks (HumanEval, GSM8k, etc.) vrijwel gelijk aan het baseline-model.
  • Voor Qwen3-30B-A3B leidt pure router-based speculatie tot een daling in nauwkeurigheid (vooral bij wiskundige taken), maar de Hybrid-PF strategie (gebruik van de estimator in vroege lagen) herstelt het grootste deel van deze nauwkeurigheid (bijv. 37% van de gap op GSM8k).
• Hit Rates: De quasi-hidden state bereikt een hit rate van ~90% in lagen met lage drift. De neurale estimator verhoogt dit in de vroege lagen van Qwen3-30B-A3B met ongeveer 25%.

Betekenis

Dit werk maakt het haalbaar om grote, open-source MoE-modellen lokaal te draaien op hardware met beperkt VRAM (zoals consumenten-GPU's), zonder dat de inferentiesnelheid onacceptabel laag wordt door CPU-GPU-transfers.

Door de "I/O-bound" aard van MoE-inferentie om te zetten in een meer "compute-bound" proces door overlapping, wordt de praktische bruikbaarheid van deze modellen voor lokale toepassingen aanzienlijk vergroot. Het paper biedt bovendien een nieuwe richting voor "router-free inference", waarbij de router in inferentietijd kan worden vervangen door een snellere estimator, wat de parallelisatie verder kan verbeteren.