EAGLE-Pangu: Accelerator-Safe Tree Speculative Decoding on Ascend NPUs

Dit paper introduceert EAGLE-Pangu, een reproduceerbaar systeem dat boomgestructureerde speculatieve decoding op Ascend NPUs veilig implementeert voor Pangu-modellen, waardoor de doorvoersnelheid aanzienlijk wordt verbeterd ten opzichte van traditionele methoden.

De kern

Stel je voor dat een grote taalmodel (zoals een slimme chatbot) een verhaal schrijft. Normaal gesproken werkt deze bot als een zeer voorzichtig, maar traag schrijver: hij bedenkt één woord, checkt of dat woord klopt, schrijft het op, bedenkt het volgende woord, checkt dat weer, en zo verder. Dit is als het bouwen van een muur steen voor steen, waarbij je bij elke steen even stopt om te kijken of de mason (de meesterbouwer) het goed vindt. Dit kost veel tijd, vooral als de muur lang moet worden.

Het probleem:
De "meesterbouwer" (het grote model) is heel slim, maar ook erg traag. De "snelle schrijver" (een kleiner model) kan veel sneller woorden bedenken, maar maakt soms fouten.

De oplossing (Speculatieve Decoding):
Om sneller te zijn, laten we de snelle schrijver eerst een heel stuk van de muur opwerpen (bijvoorbeeld 5 woorden tegelijk). De meesterbouwer kijkt dan niet naar één steen, maar naar die hele hoop en zegt: "Ja, de eerste drie zijn perfect, de vierde is een beetje scheef, en de vijfde is helemaal verkeerd." De bot accepteert dan de eerste drie en begint pas weer bij de vierde. Dit is veel sneller!

Het nieuwe probleem (De boom):
Soms is het zelfs nog slimmer om niet één lijn te tekenen, maar een boom te maken. De snelle schrijver bedenkt: "Wat als we woord A zeggen? Of misschien woord B? Of C?" Hij maakt een takkenstructuur. De meesterbouwer moet dan al die takken tegelijk controleren.

Hier komt het papier EAGLE-PANGU om de hoek kijken. Ze hebben deze slimme "boom-methode" geprobeerd op een heel specifieke, krachtige computerchip (de Ascend NPU van Huawei). Maar dat was als proberen een Formule 1-auto op een fietspad te rijden: het systeem was te fragiel. De chip verwerkte de "boom" niet goed, de adressen van de woorden waren soms negatief (wat de chip niet mag), en de geheugenopslag raakte in de war.

Wat heeft EAGLE-PANGU gedaan? (De analogieën)

1. De "Veilige Boom" (Accelerator-Safe Tree Tensorization):
   Stel je voor dat je een boom tekent op een bordje dat alleen positieve getallen accepteert. Als je een tak hebt die naar "geen ouder" wijst (de stam), zou je normaal -1 gebruiken. Maar de computerchip schreeuwt: "Fout! Geen negatieve getallen!"
   EAGLE-PANGU lost dit op door een dummy-stam toe te voegen. In plaats van te zeggen "geen ouder", zeggen ze "ouder is de stam (nummer 0)". Zo weet de computer altijd waar hij moet kijken, zonder in paniek te raken. Het is alsof je een extra, onzichtbare steunpaal in de grond zet zodat elke tak ergens aan vastzit.

2. De "Slimme Opbergkast" (Branchable KV-cache):
   Tijdens het schrijven moet de bot alle eerdere woorden onthouden in een soort geheugenkast. Als de snelle schrijver 5 takken bedenkt, moet de bot 5 versies van die kast hebben, zonder dat ze elkaar verstoren.
   EAGLE-PANGU heeft een systeem bedacht dat als een magische fotokopieerapparaat. Het maakt snel kopieën van de geheugenkast voor elke tak, zodat de meesterbouwer ze allemaal tegelijk kan checken. Zodra de meesterbouwer zegt "Tak 1 is goed!", wordt die specifieke kopie direct de nieuwe "echte" kast, en de andere worden weggegooid. Dit gaat razendsnel en voorkomt dat de bot in de war raakt over welke woorden hij al heeft geschreven.

3. De "Verkeersregelaar" (Tree Attention Masking):
   In een boom mogen de takken elkaar niet "horen" als ze niet met elkaar verbonden zijn. Als tak A en tak B naast elkaar staan, mag A niet weten wat B denkt, tenzij ze van dezelfde ouder komen.
   EAGLE-PANGU bouwt een slimme verkeersregelaar die precies regelt wie naar wie mag kijken. Dit zorgt ervoor dat de computer niet per ongeluk informatie van de verkeerde tak gebruikt, wat de antwoorden zou verpesten.

Het resultaat:
Door deze drie dingen slim te regelen, werkt de boom-methode nu perfect op deze specifieke chips.

• Snelheid: De bot is gemiddeld 1,27 keer sneller.
• Bij piekmomenten: Soms is hij zelfs 2,46 keer sneller!

Conclusie:
Dit papier is niet zozeer een nieuwe manier van "denken" voor de AI, maar een technische handleiding om een bestaande, slimme methode (het tekenen van een boom van woorden) veilig en snel te laten werken op een specifieke, krachtige maar kieskeurige computerchip. Het is alsof ze een Formule 1-auto hebben aangepast zodat hij niet meer vastloopt op het fietspad, maar er juist razendsnel overheen kan racen.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "EAGLE-Pangu: Accelerator-Safe Tree Speculative Decoding on Ascend NPUs", geschreven in het Nederlands.

1. Het Probleem

De belangrijkste bottleneck bij het serveren van Large Language Models (LLM) is de autoregressieve decoding, waarbij elke nieuwe token een nieuwe forward-pass door een groot "teacher"-model vereist. Dit beperkt de doorvoer en verhoogt de latentie.

Hoewel speculatieve decoding (het gebruik van een kleiner "draft"-model om kandidaten te genereren die door de teacher worden geverifieerd) dit probleem vermindert, is de implementatie van boomgestructureerde speculatieve decoding (waarbij meerdere continuaties parallel worden geëvalueerd) vaak fragiel bij porting naar heterogene backends, zoals de Ascend NPUs van Huawei.

De specifieke uitdagingen op deze hardware zijn:

• Niet-uitwisselbare semantiek: Verschillen in KV-cache lay-outs, attention masking en indexering tussen frameworks.
• Ongedefinieerde indexering: Op Ascend NPUs kunnen negatieve indices (vaak gebruikt als 'sentinel' waarden voor de root van een boom) ongedefinieerd gedrag of runtime-fouten veroorzaken.
• Fused Kernels: Geoptimaliseerde attention-kernels stellen strengere eisen aan maskervormen en randvoorwaarden dan "eager" implementaties.
• Foutgevoeligheid: Een naïeve port kan leiden tot kwaliteitsverlies, sporadische crashes of irreproduceerbare prestaties, zelfs als de code "draait".

2. Methodologie: EAGLE-PANGU

Het paper introduceert EAGLE-PANGU, een reproduceerbaar systeem dat de EAGLE-3 boom-speculatieve decoding portt naar een Pangu teacher-backend op Ascend NPUs. De aanpak focust op systeemniveau-fouten en biedt abstracties die de decoding-semantiek behouden terwijl ze voldoen aan de accelerator-beperkingen.

De kernmethodologie bestaat uit drie pijlers:

A. Branchable KV-cache Abstractie

Om de staat van de geverifieerde prefix te scheiden van de speculatieve takken, wordt een cache-manager gebouwd op de HuggingFace Cache API.

• Isolatie: Bij elke iteratie worden geïsoleerde branch-caches gekloond (Replicate) van de gecommitteerde cache. Dit voorkomt dat het uitbreiden van één tak de staat van andere takken of de hoofdcache corrumpeert.
• Commit-modes: Er zijn twee strategieën om de cache te updaten na acceptatie:
  1. Length-based: Behoud de originele prefix en voeg de eerste $A$ stappen van de gekozen tak toe.
  2. Path-index-based: Herordent de cache op basis van een index-mapping. Voor prestaties wordt een "fast reorder" gebruikt die alleen het nieuwe segment herschikt als de prefix-ordening behouden blijft.

B. Accelerator-Veilige Boom-Tensor Semantiek

Om problemen met negatieve indices en out-of-bounds toegang op te lossen, wordt een speciaal indexeringsschema ontworpen:

• Dummy-Root Indexering: In plaats van een sentinel-waarde (zoals -1) voor de root-ouder, wordt een dummy-rij op index 0 toegevoegd aan de tensor. Alle ouder-indices worden verschoven zodat ze binnen het bereik $[0, M]$ vallen.
• Ancestor Tabellen: Voor operaties die meerdere voorouders nodig hebben, worden tabellen geconstrueerd die gegarandeerd veilige indices genereren.
• Invariant Checks: Voordat fused kernels worden gestart, worden lichte checks uitgevoerd om cycli, ongeldige bereiken en geldigheids-sluiting te verifiëren. Dit voorkomt stille corruptie.

C. Fused Tree-Masked Teacher Executie

De teacher-evaluatie wordt uitgevoerd als één gebatchte forward-pass over alle speculatieve tokens.

• Boom-Attention Mask: Een 4D-maskeertensor wordt gegenereerd die toestaat dat token $k$ alleen naar token $j$ attendeert als $j$ een voorouder is van $k$ in de boomstructuur. Dit voorkomt informatielekken tussen takken.
• Dual-Mode Protocol: Het systeem ondersteunt twee modi:
  1. Reference Mode: Schakelt fused attention uit voor debuggen en invariant-checks.
  2. Performance Mode: Schakelt fused attention in voor maximale doorvoer.

3. Belangrijkste Bijdragen

1. Branchable KV-cache Manager: Een interface die commit- en speculatieve staten strikt scheidt, wat correcte cache-kloning en updates mogelijk maakt zonder backend-specifieke kennis.
2. Veilige Tensor Semantiek: Een ontwerp dat negatieve padding-indexen vervangt door geldige dummy-indices, waardoor ongedefinieerd gedrag op Ascend NPUs wordt geëlimineerd.
3. Fused-Kernel Compatibiliteit: Een geïntegreerde teacher-verification path die boom-masks ondersteunt voor fused attention, met een fallback naar eager execution voor debugging.
4. Reproduceerbaarheid: Een gedistribueerde pipeline met gestructureerde traces en deterministische sharding, wat ablatie-studies en debugging over verschillende uitvoeringsmodi mogelijk maakt.

4. Resultaten

Het systeem werd geëvalueerd op 240 turns (van MT-Bench en HumanEval-stijl prompts) op Ascend NPUs.

• Doorvoerverbetering: In vergelijking met alleen greedy decoding (teacher-only) bereikte EAGLE-PANGU een gemiddelde toename in end-to-end doorvoer van 1.27x.
• Tail-prestaties: De verbetering is het grootst in de lange staart, met een snelheidswinst van 2.46x bij het 99e percentiel (p99).
• Acceptatie-lengte: De gemiddelde geaccepteerde lengte per verificatiestap was 3.17 tokens. Er is een sterke correlatie tussen de geaccepteerde lengte en de snelheidswinst.
• Budget Sensitiviteit: Er is een "sweet spot" gevonden voor de boomgrootte (bijv. 16 knopen met een diepte van 10). Grotere budgets leiden tot afnemende meeropbrengsten door overhead bij het construeren van masks en het herschikken van caches.
• Negatief Resultaat: Het simplistisch afkappen van de context van het draft-model (fixed-window truncation) bleek schadelijk. Dit verlaagde de acceptatiekans aanzienlijk en keerde de snelheidswinst om, omdat het draft-model vaak afhankelijk is van verre geschiedenis (zoals aangetoond door attention-profielen).

5. Significantie

Dit paper is significant omdat het een praktische, reproduceerbare route biedt om geavanceerde boom-speculatieve decoding te implementeren op Ascend NPUs, een hardwareplatform dat vaak als "zwart doos" wordt ervaren door complexe kernel-eisen.

• Correctheid boven snelheid: Het paper benadrukt dat correctheid (via veilige indexering en invariant checks) de voorwaarde is voor snelheid op deze hardware.
• Hardware-onafhankelijkheid: Door abstracties zoals de branchable cache en veilige tensor mapping, maakt het systeem de techniek minder fragiel bij porting naar verschillende accelerators.
• Productie-klaar: De implementatie bevat mechanismen voor debugging, tracing en distributie, wat essentieel is voor de adoptie in productiesystemen.

Samenvattend toont EAGLE-PANGU aan dat boom-speculatieve decoding de doorvoer van LLM's op Ascend NPUs aanzienlijk kan verbeteren zonder ingrijpende herschrijving van kernels, mits de systeemarchitectuur rekening houdt met de specifieke beperkingen van de accelerator.