RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators

RedFuser is een automatisch framework dat cascaderende reductie-operaties in AI-modellen efficiënt samenvoegt tot een enkele lus, waardoor een snelheidsverhoging van 2 tot 5 keer wordt bereikt ten opzichte van bestaande compilers.

De kern

RedFuser: De "Super-Kok" voor AI-keukens

Stel je voor dat een AI-model (zoals een slimme chatbot) een gigantische receptenboek is. Om een antwoord te geven, moet de computer duizenden kleine berekeningen uitvoeren. In de wereld van AI noemen we deze berekeningen "operatoren".

Het probleem is dat de huidige software (de "koks") deze berekeningen vaak stap voor stap, één voor één, uitvoert. Het is alsof je een taart maakt, maar eerst deeg maakt, dat in een bak doet, de bak uit de oven haalt, deeg weer uit de bak haalt, het op een bord legt, en pas dan deeg weer in de oven doet voor de volgende laag. Je loopt constant heen en weer naar de oven en de koelkast. Dit kost veel tijd en energie.

RedFuser is een nieuw slimme tool die dit probleem oplost. Het is een automatisch systeem dat deze losse stappen samenvoegt tot één efficiënte, soepele beweging.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Kettingreactie" van Wachten

In AI-modellen zijn er bepaalde berekeningen die in een ketting hangen. Je kunt pas de tweede stap doen als de eerste klaar is.

• Voorbeeld: Stel je voor dat je een grote menigte mensen moet tellen, en daarna het gemiddelde gewicht van die mensen moet berekenen.
• De oude manier: Je telt eerst iedereen (stap 1), schrijft het getal op een briefje, loopt naar een ander kantoor, leest het briefje, en begint dan pas met wegen (stap 2). Tussendoor loop je heen en weer, en je moet het briefje steeds opnieuw lezen.
• De inefficiëntie: De computer moet steeds dezelfde gegevens uit het geheugen (de koelkast) halen, verwerken, opslaan, en weer ophalen. Dit is als een bezorgdienst die elke bestelling apart aflevert in plaats van één grote vrachtwagen te vullen.

2. De Oplossing: RedFuser als de "Meester-Kok"

RedFuser kijkt naar deze losse stappen en zegt: "Wacht even, jullie doen dit allemaal in dezelfde keuken. Laten we het in één keer doen!"

Het doet dit door twee slimme trucjes toe te passen:

Truc 1: De "Alles-in-Één" Pan (Operator Fusion)

In plaats van het deeg te maken, op te bergen, en later weer te gebruiken, doet de kok het deeg maken, het bakken en het serveren in één continue stroom.

• In de computer: RedFuser pakt de losse berekeningen en smelt ze samen tot één enkele, krachtige instructie. De computer hoeft de gegevens niet meer op te slaan en weer op te halen. Het blijft in het snelle, lokale geheugen van de chip.
• Het resultaat: Geen meer heen en weer lopen naar de koelkast. Alles gebeurt in één beweging.

Truc 2: De "Streaming" Methode (Incremental Computation)

Soms is de menigte zo groot dat hij niet in één keer in de keuken past. De oude methoden probeerden alles in het geheugen te proppen, wat leidde tot chaos en vertraging.

• De nieuwe manier: RedFuser gebruikt een "streaming" aanpak. Stel je voor dat je een rivier hebt. In plaats van te wachten tot de hele rivier is opgevangen in een reservoir (wat te groot is), vang je het water op terwijl het stroomt, verwerk je het direct, en stuur je het direct door.
• Waarom is dit slim? Je hoeft niet te wachten tot alles binnen is. Je kunt alvast beginnen met het volgende stukje terwijl je het vorige verwerkt. Dit betekent dat de computer nooit stopt om te wachten en altijd voluit draait, zelfs bij gigantische hoeveelheden data.

3. Wat levert dit op?

De auteurs van het paper hebben RedFuser getest op moderne AI-chips (zoals die van NVIDIA). De resultaten zijn indrukwekkend:

• 2 tot 5 keer sneller: AI-modellen draaien veel sneller dan met de huidige standaardsoftware.
• Net zo goed als handwerk: Vaak schrijven experts hun eigen code om dit snel te maken. RedFuser doet dit automatisch en komt uit op hetzelfde hoge niveau, maar dan zonder dat een mens urenlang hoeft te puzzelen.
• Toepasbaar op alles: Of het nu gaat om het begrijpen van taal (zoals in chatbots), het aanbevelen van films, of het analyseren van medische beelden; waar er deze "ketting-berekeningen" zijn, werkt RedFuser.

Samenvattend

Stel je voor dat AI-berekeningen een drukke metrostation zijn.

• De oude software laat elke passagier een apart ticket halen, wachten in een rij, en dan pas de trein opstappen.
• RedFuser is het systeem dat alle passagiers in één groep samenbrengt, hen direct door de poort laat lopen en ze direct in de trein zet, terwijl de trein al vertrekt.

Het is een automatische "snelheidsverhoging" die ervoor zorgt dat onze AI-toekomst niet vastloopt in trage berekeningen, maar soepel en razendsnel vooruit gaat.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators" in het Nederlands.

Probleemstelling

In de huidige AI-compiler-ecosystemen (zoals TVM, PyTorch Dynamo, XLA) is operator-fusie een cruciale techniek voor prestatieoptimalisatie. Echter, deze systemen hebben moeite met het efficiënt optimaliseren van gekaskedeerde reducties (cascaded reductions).

• Definitie: Gekaskedeerde reducties zijn sequenties van reductie-operaties (zoals max, sum, argmax) waarbij er een strikte data-afhankelijkheid bestaat tussen de stappen. Een typisch voorbeeld is de safe softmax in attention-mechanismen (eerst max vinden, dan exp berekenen, dan sum nemen, en tenslotte normaliseren), gevolgd door een GEMM-operatie.
• Huidige beperkingen: Bestaande compilers behandelen deze operaties vaak als aparte loops. Dit leidt tot:
  1. Redundante geheugentoegang: Elke reductiestap laadt de invoerdata opnieuw en schrijft tussenresultaten naar het geheugen, wat de bandbreedte verzadigt.
  2. Beperkte parallelisatie: Vanwege de data-afhankelijkheid kunnen de reductiebomen niet parallel worden uitgevoerd, wat de latentie verhoogt.
  3. Cache-druk: Tussenresultaten moeten volledig in het on-chip geheugen (shared memory) worden bewaard, wat de maximale lengte van de reductie beperkt.
• Handmatige optimalisatie: Hoewel handgeschreven kernels (zoals FlashAttention) uitstekende prestaties leveren, zijn ze specifiek voor bepaalde patronen, moeilijk te generaliseren en vereisen ze expert-kennis van GPU-programmering.

Methodologie

De auteurs introduceren een formele theoretische methode om gekaskedeerde reducties te analyseren en te fuseren tot een enkele, geoptimaliseerde loop.

1. Formele Definitie en Fusie-voorwaarden
De auteurs definiëren een reeks van $I$ reductie-operaties. Om deze te kunnen fuseren, moet de $i$-de reductiefunctie $F_i$ voldoen aan specifieke algebraïsche eigenschappen:

• Decomposeerbaarheid: De functie moet kunnen worden opgesplitst in een deel dat alleen van de invoer afhangt ($G_i$) en een deel dat alleen van de afhankelijke resultaten afhangt ($H_i$), gekoppeld via een commutatief monoid-operatorsysteem ($\otimes_i$).
• Distributiviteit: De reductie-operator ($\oplus_i$) moet distributief zijn over de koppel-operator ($\otimes_i$).
• Inverteerbaarheid: De functie $H_i$ moet inverteerbaar zijn (of er moet een correctiemechanisme worden toegepast).

2. Gecombineerde Expressie en Incrementele Berekening
Op basis van deze voorwaarden leiden de auteurs een gefuseerde expressie af die de data-afhankelijkheid doorbreekt. In plaats van te wachten op het eindresultaat van de vorige reductieboom, wordt de berekening herschikt zodat invoerdata slechts één keer wordt geladen.

• Incrementele Berekening: Om het probleem van onvoldoende on-chip cache voor lange sequenties op te lossen, introduceren ze een incrementele berekeningsvorm. Hierbij worden tussenresultaten direct bijgewerkt zodra nieuwe data beschikbaar komt, in plaats van het wachten op de volledige vorige reductie. Dit verlaagt de geheugencomplexiteit van $O(L)$ naar $O(1)$, waardoor zeer lange reducties mogelijk zijn binnen de beperkte shared memory van een GPU-block.

3. Het RedFuser Framework
RedFuser is een tweestaps-framework gebouwd op de TVM compiler-stack:

• Stap 1: Symbolische Afleiding: Een engine analyseert de abstract syntax tree (AST) van de computationele grafiek. Het gebruikt een "Fixed-Point Analysis" om te bepalen of een patroon voldoet aan de fusievoorwaarden en genereert automatisch de symbolische uitdrukkingen voor de gefuseerde en incrementele berekening.
• Stap 2: Hardware-bewuste Codegeneratie: De gefuseerde expressies worden vertaald naar een tile-gebaseerde representatie (TensorIR). RedFuser genereert twee strategieën:
  • Single-Segment: De hele reductie in één CTA (Cooperative Thread Array) met incrementele updates.
  • Multi-Segment: Parallelle verwerking van segmenten door meerdere CTAs, gevolgd door een samenvoeging (reductie) van de partiële resultaten.
  • De code wordt vervolgens geoptimaliseerd met TileLang voor specifieke GPU-architecturen (bijv. gebruik van TMA, WGMMA op Hopper, cp.async op Ampere).

Belangrijkste Bijdragen

1. Formele Methodologie voor Gekaskedeerde Reducties: Het eerste werk dat een wiskundige basis legt voor het analyseren en fuseren van complexe, data-afhankelijke reductieketens in een enkele loop.
2. Incrementele Berekeningsvorm: Een innovatieve aanpak die de cache-beperkingen van GPU's omzeilt, waardoor fusie mogelijk is voor willekeurige reductielengtes zonder verlies van parallelisme.
3. RedFuser Framework: Een volledig geautomatiseerd systeem dat deze theorie implementeert, ondersteunend voor diverse werklasten (Attention, MoE, Quantization) zonder handmatige kernel-ontwikkeling.
4. State-of-the-Art Prestaties: Het framework levert kernels op die vergelijkbaar zijn met of zelfs beter presteren dan handgeschreven optimalisaties (zoals FlashAttention), maar dan gegenereerd door een compiler.

Resultaten

De auteurs hebben RedFuser geëvalueerd op moderne GPU's (NVIDIA A10 en H800) tegenover bestaande compilers (PyTorch Eager, PyTorch Dynamo/Inductor, TVM) en hand-geoptimaliseerde bibliotheken.

• Snelheidswinst: RedFuser bereikte een snelheidswinst van 2x tot 5x ten opzichte van state-of-the-art AI compilers.
• Vergelijking met Handgeschreven Kernels:
  • Voor Multi-Head Attention (MHA) en Multi-Latent Attention (MLA) presteerde RedFuser gelijk aan of beter dan FlashAttention2 en FlashMLA (bijv. 102% van FlashMLA prestaties).
  • Voor MoE-routing en FP8 Quant + GEMM (waarvoor geen standaard hand-geoptimaliseerde kernels bestaan) werden enorme verbeteringen gezien (tot 12.1x sneller dan TVM).
• Flexibiliteit: De incrementele modus bleek essentieel voor lange sequenties (zoals bij LLM-decoding), waar de niet-incrementele modus faalde vanwege geheugenbeperkingen. De incrementele modus maakte het mogelijk om optimale parallelisme-configuraties te bereiken die anders ontoegankelijk waren.

Betekenis en Impact

Dit werk is significant omdat het de kloof overbrugt tussen handmatige, hooggespecialiseerde optimalisaties en generieke compiler-automatisering.

• Het toont aan dat complexe, data-afhankelijke patronen in AI-modellen (die vaak de bottleneck vormen) volledig geautomatiseerd kunnen worden geoptimaliseerd.
• Het vermindert de afhankelijkheid van GPU-experts voor het schrijven van snelle kernels voor nieuwe architecturale patronen (zoals MoE of FP8-quantisatie).
• Het biedt een schaalbare oplossing voor de groeiende complexiteit van AI-modellen, waarbij de efficiëntie van training en inferentie kritiek is voor kosteneffectiviteit en snelheid.

Kortom, RedFuser transformeert een specifiek, moeilijk optimaliseerbaar probleem in een gestandaardiseerd, automatisch oplosbaar proces, terwijl het de prestaties van de beste handmatige implementaties behoudt of overtreft.