Batch Speculative Decoding Done Right

Titel: Batch Speculative Decoding Done Right

Auteurs: Ranran Haoran Zhang, Soumik Dey, Ashirbad Mishra, Hansi Wu, Binbin Li, Rui Zhang.

---

1. Het Probleem: Output Equivalency en het "Ragged Tensor"-Probleem

Speculatieve decoding (Speculative Decoding) is een techniek om de inferentie van Large Language Models (LLM's) te versnellen door een klein "draft"-model tokens te laten voorspellen die parallel worden geverifieerd door het grotere "target"-model. De fundamentele eis voor een geldige implementatie is output-equivalentie: de verdeling van de gegenereerde tokens moet exact identiek zijn aan die van standaard autoregressieve generatie.

Het paper identificeert een kritiek falen in bestaande batch-implementaties (zoals BSP en DSD):

• Output Corruptie: Bestaande methoden produceren vaak "gibberish", herhalende tokens of <unk>-symbolen in plaats van correcte tekst, zelfs als ze hoge doorvoer (throughput) claimen.
• Oorzaak (Ragged Tensor Probleem): Wanneer meerdere sequenties in één batch worden verwerkt, accepteren ze een verschillend aantal gedraaide (draft) tokens per verificatiestap. Dit leidt tot "ragged" (onregelmatige) tensors.
• Desynchronisatie: Bestaande implementaties handelen deze onregelmatigheid verkeerd af, wat resulteert in een desynchronisatie van:
  • Position IDs: De posities van tokens komen niet meer overeen met de attention-mechanismen (zoals RoPE).
  • Attention Masks: De maskers verwijzen naar verkeerde tokens.
  • KV-Cache: De Key-Value cache bevat tokens die niet corresponderen met de huidige sequentielengte of positie.

De auteurs tonen aan dat bestaande methoden (Masking, Rollback, Dynamic Padding) ofwel de output-equivalentie schenden of inefficiënt zijn door tokens weg te gooien (rollback) of onnodige padding toe te voegen zonder correcte hersynchronisatie.

---

2. Methodologie

De auteurs introduceren een nieuw raamwerk dat de synchronisatie-invarianten formeel definieert en twee algoritmen presenteert: EQSPEC en EXSPEC.

A. Formele Invarianten

Voor geldige batch-speculatieve decoding moeten drie invarianten worden gehandhaafd na elke verificatiestap:

1. Rechthoekige Alignering (Rectangular Alignment): Alle sequenties in de batch moeten na padding dezelfde totale lengte hebben ($p_i + c_i = L$).
2. Contiguïteit van Position IDs: Position IDs moeten continu zijn en correct corresponderen met de content-tokens (niet met padding), essentieel voor positi-afhankelijke attention.
3. KV-Cache Synchronisatie: De cache moet exact overeenkomen met de geaccepteerde token-stream.

B. EQSPEC: Het Correcte Algoritme

EQSPEC is het eerste algoritme dat output-equivalentie garandeert door expliciete synchronisatie.

• Mechanisme: Na elke verificatiestap voert EQSPEC een "Unpad-Append-Repad" procedure uit:
  1. Unpad: Verwijder bestaande padding.
  2. Append: Voeg de geaccepteerde tokens en de "bonus token" (van het target-model) toe.
  3. Repad: Voeg nieuwe padding toe aan de linkerkant zodat alle sequenties weer dezelfde lengte hebben.
  4. Hersynchronisatie: Herbereken de position IDs, update de attention masks en verschuif de KV-cache entries om de nieuwe posities te matchen.
• Kosten: Deze hersynchronisatie is noodzakelijk voor correctheid, maar introduceert een aanzienlijke overhead (tot 40% van de rekentijd bij batchgrootte 8).

C. EXSPEC: Overheadreductie via Cross-Batch Scheduling

Om de hoge overhead van EQSPEC te verminderen, introduceert EXSPEC een Cross-Batch Scheduling strategie.

• Pool-benadering: In plaats van vaste batches te forceren, houdt EXSPEC een pool van actieve sequenties bij.
• Dynamische Groepering: De scheduler probeert sequenties met identieke lengte te groeperen in batches.
• Voordeel: Als alle sequenties in een batch dezelfde lengte hebben en hetzelfde aantal tokens accepteren, is er geen padding nodig en vervalt de hersynchronisatiekosten volledig ($\delta_i = 0$).
• Fallback: Alleen wanneer groepering faalt (ongelijke lengtes), wordt de duurdere "Unpad-Append-Repad" methode van EQSPEC gebruikt.

---

3. Belangrijkste Bijdragen

1. Identificatie van Fundamentele Fouten: Het paper bewijst dat alle bestaande batch-speculatieve decoding-implementaties (BSP, DSD) output-equivalentie schenden door het "ragged tensor"-probleem verkeerd op te lossen, wat leidt tot corrupte output.
2. Formalisatie van Synchronisatie: De auteurs definiëren de strikte invarianten (Position IDs, Attention Masks, KV-Cache) die nodig zijn voor correctheid.
3. EQSPEC: Het eerste algoritme dat wiskundig bewijst dat output-equivalentie behouden blijft door expliciete hersynchronisatie na elke stap.
4. EXSPEC: Een geavanceerde scheduler die de overhead van hersynchronisatie minimaliseert door sequenties met dezelfde lengte dynamisch te groeperen, waardoor de schaalbaarheid verbetert.
5. Analyse van Kosten: Het paper toont aan dat de aligneringsoverhead superlineair groeit met de batchgrootte en dat dit een inherente kostenpost is voor correctheid, geen implementatiefout.

---

4. Resultaten

De methoden zijn geëvalueerd op SpecBench met modelparen zoals Vicuna-7B/68M, Qwen3-8B/0.6B en GLM-4-9B/0.6B.

• Output Equivalency (Correctheid):
  • EQSPEC en EXSPEC: Bereiken ~95% exacte match met standaard autoregressieve decoding. De resterende 5% divergentie wordt toegeschreven aan numerieke non-determinisme in GPU-inferentie (floating-point), niet aan algoritmische fouten.
  • Bestaande methoden (BSP, DSD): Falen catastrofisch met near-zero exacte match (bijv. BSP produceert herhalingen, DSD produceert <unk>).
• Doorvoer (Throughput):
  • Bij een batchgrootte van 8 bereikt EXSPEC een 3x snelheidswinst ten opzichte van batchgrootte 1, terwijl het de correctheid behoudt.
  • EQSPEC toont negatieve schaling bij zeer grote batches (>8) door de toenemende hersynchronisatiekosten, wat de noodzaak van EXSPEC onderstreept.
• Overhead Analyse:
  • Hersynchronisatie (alignment) consumeert tot 40% van de rekentijd bij batchgrootte 8 en stijgt naar ~47% bij batchgrootte 16.
  • EXSPEC reduceert deze overhead aanzienlijk door homogene groepen te vormen, wat de effectiviteit van cross-batch scheduling bewijst.

---

5. Betekenis en Impact

• Fundamentele Correctheid: Dit paper stelt een nieuwe standaard voor batch-speculatieve decoding. Het toont aan dat snelheid ten koste van correctheid (output-equivalentie) geen optie is; een snelle maar foutieve generator is nutteloos voor productie.
• Productie-Ready: De implementatie biedt een betrouwbare basis voor het versnellen van LLM-inferentie in productieomgevingen zonder de kwaliteit van de output te compromitteren.
• Inzicht in Kosten: Het paper onthult dat de kosten van het handhaven van correctheid bij onregelmatige batches (ragged tensors) een fundamentele barrière zijn die alleen kan worden verzacht (niet geëlimineerd) door slimme scheduling zoals EXSPEC.
• Open Source: De auteurs hebben de code openbaar gemaakt, waardoor de gemeenschap een geverifieerde referentie-implementatie heeft voor correcte batch-speculatieve decoding.

Kortom, dit werk corrigeert een breed verspreid misverstand in de LLM-inferentie-gemeenschap en levert de eerste wiskundig correcte en schaalbare oplossing voor batch-speculatieve decoding.