💻 computer science

Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variants

Dit artikel introduceert Flashlight, een PyTorch-native compilerframework dat automatisch gefuseerde, FlashAttention-achtige kernels genereert voor willekeurige en data-afhankelijke attentiepatronen zonder te vertrouwen op statische sjablonen, waardoor het superieure flexibiliteit en concurrerende prestaties biedt ten opzichte van bestaande oplossingen zoals FlexAttention.

Oorspronkelijke auteurs: Bozhi You, Irene Wang, Zelal Su Mustafaoglu, Abhinav Jangda, Angélica Moreira, Roshan Dathathri, Divya Mahajan, Keshav Pingali

Gepubliceerd 2026-05-22

📖 6 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Bozhi You, Irene Wang, Zelal Su Mustafaoglu, Abhinav Jangda, Angélica Moreira, Roshan Dathathri, Divya Mahajan, Keshav Pingali

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Het Grote Plaatje: De "Stuip" in AI

Stel je voor dat een Large Language Model (zoals degene waarmee je praat) een enorme fabriek is die probeert een complex product samen te stellen (een zin of een eiwitstructuur). Het meest kritieke deel van deze fabriek is de "Attention"-afdeling. Hier bepaalt de machine welke stukjes informatie belangrijk zijn en welke genegeerd moeten worden.

Het probleem? De huidige manier waarop deze afdeling werkt, is als een verkeersopstopping.

De Oude Manier: De fabrieksarbeiders (de computercode) moeten stoppen, een notitie op een stuk papier schrijven (gegevens opslaan in trage geheugen), naar een ander station lopen, de notitie oppakken en vervolgens weer aan het werk gaan. Dit gebeurt keer op keer. Het is traag en verspillen energie.
De "Flash"-Oplossing: Een paar jaar geleden bedachten ingenieurs FlashAttention. Ze ontdekten hoe ze de arbeiders alle notities in hun zakken konden houden (snel geheugen) en het hele werk in één continue sprint konden doen. Dit was een enorme snelheidswinst, maar het was als een gespecialiseerde raceauto: het werkte alleen op één specifiek type circuit (standaard attention). Als je een ander type auto wilde rijden (een nieuwe, experimentele attention-methode), moest je een hele nieuwe raceauto van hand tot hand bouwen.

Het Probleem: Te Veel Aangepaste Raceauto's

Recentelijk hebben wetenschappers vele nieuwe, fancy manieren bedacht om "Attention" te doen om AI slimmer of efficiënter te maken (zoals "Differential Attention" of "Gated Self-Attention").

De Knelpunt: Om deze nieuwe methoden snel te maken, heb je meestal een menselijk expert nodig die voor elk een aangepaste "raceauto" (een gespecialiseerde computerkernel) met de hand codeert. Dit duurt eeuwen.
Het Midden (FlexAttention): Een tool genaamd FlexAttention probeerde dit op te lossen door bouwers een set Lego-sjablonen te geven. Als je nieuwe auto in het sjabloon paste, werkte het snel. Maar als je auto raar of uniek was (niet in het sjabloon paste), kon FlexAttention je niet helpen, en was je terug bij het trage, handmatige coderen.

De Oplossing: FLASHLIGHT (De Universele Fabrieksupgrade)

Hier komt FLASHLIGHT. De auteurs beschrijven het als een compiler-natief framework. In eenvoudige termen is het een slimme "autopilot" die direct in de PyTorch-software is ingebouwd (de taal die de meeste AI-onderzoekers gebruiken).

Zo werkt FLASHLIGHT, met een analogie:

1. De "Slimme Chef" versus het "Receptenboek"

Oude Manier: Je hebt een kookboek (de code). Als je een standaard soep wilt maken, zegt het kookboek je om te hakken, te koken en te serveren. Als je een rare soep wilt, moet je een chef inhuren om een nieuw recept te bedenken en een nieuwe pagina in het kookboek te schrijven.
FlexAttention: Het kookboek heeft een "Speciale Soep Sectie". Als je soep in het "Speciale Soep"-sjabloon past, is het snel. Zo niet, dan zit je vast.
FLASHLIGHT: FLASHLIGHT is als een superslimme keukenrobot die kijkt hoe je kookt. Het maakt niet uit of je een standaard soep of een rare, experimentele soep maakt. Het kijkt naar je acties (de code), beseft dat je veel aan het hakken en koken bent, en zegt: "Hé, ik kan al deze stappen combineren in één super-efficiënte beweging, zodat je niet heen en weer naar de koelkast hoeft te lopen."

Het herschrijft je code automatisch om even snel te zijn als een met de hand gebouwde raceauto, zonder dat je hoeft te weten hoe je de auto bouwt of in een sjabloon moet passen.

2. Hoe het de Magie doet (De Drie Trucs)

Het artikel legt drie belangrijkste "trucs" uit die FLASHLIGHT gebruikt om dingen te versnellen:

Truc A: De "Degradering" (Samenvoegen van Stappen)
Stel je voor dat je een taart bakt. Stap 1 is bloem mengen. Stap 2 is eieren toevoegen. Normaal meng je de bloem, zet je de kom neer, pak je een nieuwe kom op en voeg je eieren toe.
FLASHLIGHT zegt: "Waarom stoppen? Houd de bloem in de kom en voeg de eieren daar gewoon direct toe." Het voegt aparte stappen samen tot één continue stroom, zodat gegevens niet heen en weer hoeven te reizen.
Truc B: De "Wiskundige Magie" (Algebraïsche Transformatie)
Soms moet je, om veilig te zijn, een berekening twee keer doen (zoals de temperatuur controleren, en dan nog een keer om zeker te zijn). Dit is traag.
FLASHLIGHT gebruikt een wiskundige truc (een "homomorfisme") om te beseffen dat je beide controles tegelijkertijd kunt doen terwijl je kookt, in plaats van te stoppen om ze apart te doen. Het is alsof je de oventemperatuur controleert terwijl je het beslag roert, in plaats van te stoppen om naar de oven te kijken en dan weer terug te gaan om te roeren.
Truc C: De "Tiling" (De Truck Volpakken)
Stel je voor dat je verhuist. Als je één stoel per keer draagt, duurt het eeuwen. Als je een hele truck (een "tile") vol meubels pakt en die één keer rijdt, is het snel.
FLASHLIGHT is slim in hoe het de truck volpakt. Het berekent de perfecte grootte voor de truck op basis van wat je verplaatst. Als een meubelstuk klein is, past het het in een hoekje zodat je geen ruimte verspilt. Het zorgt ervoor dat de "truck" (geheugen) van de computer altijd vol is en efficiënt beweegt.

Wat hebben ze bewezen?

De onderzoekers testten FLASHLIGHT op krachtige computerchips (NVIDIA H100 en A100).

Voor Standaard Taken: Toen ze de "standaard" attention-methoden gebruikten die FlexAttention al kon verwerken, was FLASHLIGHT even snel of zelfs sneller.
Voor Rare/Nieuwe Taken: Toen ze attention-methoden probeerden die FlexAttention niet kon verwerken (zoals die gebruikt in AlphaFold, een beroemde AI voor eiwitvouwing), was FLASHLIGHT 5 keer sneller dan de standaard, niet-geoptimaliseerde code.
Real-World Test: Ze testten het op een echt eiwitvouwmodel (AlphaFold). Het zorgde ervoor dat het model 6% tot 9% sneller liep in totaal.

De Conclusie

FLASHLIGHT is een tool die AI-ontwikkelaars toelaat om code te schrijven op de normale, gemakkelijke manier zoals ze dat altijd hebben gedaan, maar het zet die code automatisch om in een high-speed, super-efficiënte motor.

Voorheen: Je moest kiezen tussen "Gemakkelijk te schrijven maar traag" OF "Snel maar moeilijk te schrijven".
Met FLASHLIGHT: Je krijgt "Gemakkelijk te schrijven EN Snel".

Het verwijdert de noodzaak voor experts om voor elk nieuw idee speciale snelheidswinsten met de hand te coderen, waardoor wetenschappers kunnen experimenteren met nieuwe typen AI-attention-modellen zonder zich zorgen te hoeven maken dat de code te traag loopt. Het is momenteel open-source, wat betekent dat iedereen het nu kan gebruiken als een fork van PyTorch.

Technische Samenvatting: FLASHLIGHT

Probleemstelling

Het optimaliseren van attentiemechanismen is cruciaal voor het versnellen van training en inferentie van Large Language Models (LLMs). Hoewel FlashAttention de prestaties voor standaard attentie succesvol heeft verbeterd door gebruik te maken van tegelindeling (tiling) en kernelfusie om geheugenbandbreedte en lanceeroverhead te verminderen, is het een handmatig geoptimaliseerde kernel die specifiek is voor vanille-attentie. Het kan niet worden toegepast op nieuwere, complexere attentievarianten (zoals differentiële attentie, rij-/kolomgewijze gegateerde zelfattentie in AlphaFold's Evoformer, Invariant Point Attention of Rectified Sparse Attention) zonder aanzienlijke handmatige engineeringinspanning.

Bestaande oplossingen zoals FlexAttention en FlashInfer proberen deze kloof te overbruggen door statische, hogere-orde sjablonen of just-in-time compilatie te bieden voor specifieke subsets van attentievarianten. Deze benaderingen lijden echter onder beperkte expressiviteit: ze vereisen dat programmeurs hun modellen opnieuw formuleren om te passen bij vooraf gedefinieerde sjablonen. Bijgevolg blijven veel nieuwe attentieformuleringen die niet in deze sjablonen passen beperkt tot langzamere, niet-gefuseerde PyTorch-implementaties, wat hun adoptie belemmert. Bovendien mist de standaard PyTorch-compilerstapel (TorchInductor) momenteel belangrijke optimalisaties zoals reductiefusie en complexe operatorfusie over geheugengrenzen heen, wat de prestaties van gegenereerde code voor deze varianten beperkt.

Methodologie

Het artikel introduceert FLASHLIGHT, een compiler-natief framework binnen het PyTorch-ecosysteem dat is ontworpen om automatisch geavanceerde, FlashAttention-achtige kernels te genereren voor willekeurige op attentie gebaseerde programma's. In tegenstelling tot eerdere benaderingen vertrouwt FLASHLIGHT niet op statische sjablonen of vooraf gedefinieerde kernel-specialisaties. In plaats daarvan integreert het direct in de torch.compile-workflow, waarbij het standaard PyTorch-code analyseert en transformeert om samenstellende bewerkingen te ontdekken en te fuseren.

FLASHLIGHT breidt de TorchInductor-intermediare representatie (IR) en optimalisatiepasses uit via drie primaire mechanismen:

Unificatie van Reductie-IR: FLASHLIGHT modelleert tensorcontracties (zoals General Matrix Multiplication, GEMM) als gegeneraliseerde reducties binnen een geünificeerde IR die onderscheid maakt tussen parallelle dimensies ( $p$ -dimensies) en reductiedimensies ( $r$ -dimensies). Dit stelt matrixvermenigvuldigingen in staat om volledig deel te nemen aan de end-to-end fusiemotor, waardoor de kunstmatige grens wordt verwijderd die GEMM-bewerkingen doorgaans isoleert van omringende berekeningen.
Structurele Fusie met Dimensiedegradering: Het framework introduceert een fusieregel die het mogelijk maakt dat een parallelle lus ( $p$ -dimensie) uit een producer-kernel wordt "gedegradeerd" tot een reductielus ( $r$ -dimensie) binnen een gefuseerde consumer-kernel. Deze transformatie wisselt parallelisme in voor het elimineren van de hoge latentie van materialisatie van tussentijdse tensors naar globaal geheugen, waardoor bewerkingen zoals max() in softmax direct kunnen worden gefuseerd met de voorafgaande matrixvermenigvuldiging ( $QK^\top$ ).
Semantische Fusie met Algebraïsche Transformatie: Om staat-afhankelijke bewerkingen zoals softmax te fuseren, maakt FLASHLIGHT gebruik van algebraïsche eigenschappen (specifiek homomorfismen) om multi-pass-algoritmen om te zetten in single-pass online-algoritmen. Het converteert bijvoorbeeld de stabiele softmax (die twee passes vereist: het vinden van het maximum en vervolgens het optellen van exponentiële waarden) naar een online softmax-algoritme. Dit wordt bereikt door doorlopende accumulatoren bij te houden voor het maximum en de som, en de som dynamisch te herschalen wanneer de doorlopende maximum verandert, gebruikmakend van de eigenschap $e^{x-y} = e^x / e^y$ .
Tegel-bewuste Dimensie-eliminatie en Logische Grids: Het framework verfijnt de lusuitvoering door iteratie over aaneengesloten tegels. Het introduceert logische grid-dimensies om de structuur van het probleem te ontkoppelen van fysieke GPU-hardwarebeperkingen. Dit maakt flexibele, per-dimensie tegelgroottes mogelijk en stelt de eliminatie van kleine dimensies in staat die volledig binnen een enkele tegel passen, waardoor fusiekansen worden vrijgegeven voor complexe multi-stadia-berekeningen (bijvoorbeeld het fuseren van $softmax(QK^\top)V$ ) die voorheen onmogelijk waren voor de compiler.

Belangrijkste Bijdragen

Fusie op Kernel-niveau: FLASHLIGHT implementeert een mechanisme om compatibele rekenblokken (bijvoorbeeld matmul + softmax, matmul + matmul) automatisch te fuseren wanneer parallelisme dit toelaat, waarbij prestatiewinst wordt benut zonder handmatig kernel schrijven.
Generaliteit via Compiler-gestuurde Aanpak: Het ondersteunt zowel standaard als opkomende attentiemechanismen via een geünificeerde compiler-gestuurde aanpak, waardoor de noodzaak voor statische sjablonen of aangepaste kernels wordt verwijderd.
Prestaties met Flexibiliteit: Het produceert kernels met concurrerende of superieure prestaties ten opzichte van bestaande sjabloon-gebaseerde oplossingen (zoals FlexAttention), terwijl de volledige flexibiliteit van native PyTorch-code behouden blijft.

Resultaten

De auteurs hebben FLASHLIGHT geëvalueerd op NVIDIA H100- en A100-GPU's voor diverse attentievarianten:

FlexAttention-ondersteunde Varianten: Voor varianten die worden ondersteund door FlexAttention (Vanilla, ALiBi, Softcap, Causal, Sliding Window, enz.) is de door FLASHLIGHT gegenereerde code concurrerend met of sneller dan FlexAttention. Specifiek voor score_mod-varianten is FLASHLIGHT tot 1,48× sneller vanwege eenvoudigere gefuseerde kernels die de overhead van het verwerken van blokmaskers vermijden. Voor block_mask-varianten is de kerneluitvoering van FlexAttention sneller vanwege het overslaan van sparsiteit, maar vermijdt FLASHLIGHT de aanzienlijke overhead van het construeren van sparsere blokmaskers tijdens runtime.
Complexe Varianten (Niet-ondersteund door FlexAttention): Voor attentievarianten die niet uitdrukbaar zijn in het sjabloon van FlexAttention (Differentiële Attentie, rij-/kolomgewijze gegateerde zelfattentie van AlphaFold's Evoformer), presteert FLASHLIGHT aanzienlijk beter dan standaard torch.compile.
- Differentiële Attentie: Toont hogere snelheidswinsten op H100 in vergelijking met A100.
- Evoformer: FLASHLIGHT behaalt snelheidswinsten van 5× of meer in vergelijking met torch.compile op beide GPU-architecturen.
End-to-End Inferentie:
- AlphaFold2: FLASHLIGHT verbetert de inferentielatentie met 6% tot 9% voor rij-/kolomgewijze gegateerde zelfattentielagen.
- LLaMa-3.2-1B: In een vLLM-systeem met een Mooncake-conversatietrace presteert FLASHLIGHT beter dan FlexAttention voor Softcap-attentie. Voor Causal-attentie (een blokmasker-variant) blijft FlexAttention iets performanter vanwege de geoptimaliseerde verwerking van sparsere blokmaskers, hoewel FLASHLIGHT out-of-memory-fouten voorkomt die optreden bij standaard torch.compile door materialisatie van tussentijdse tensors.

Betekenis en Claims

Het artikel beweert dat FLASHLIGHT een significante verschuiving vertegenwoordigt in hoe attentievarianten worden geoptimaliseerd, waarbij de last wordt verplaatst van handmatige engineering naar compilerautomatisering. Door het automatisch genereren van gefuseerde, FlashAttention-achtige kernels uit standaard PyTorch-code mogelijk te maken, stelt het AI-wetenschappers in staat om snel nieuwe attentiemodellen te verkennen zonder prestaties of schaalbaarheid op te offeren.

De auteurs benadrukken dat FLASHLIGHT open source is en beschikbaar is als een fork van PyTorch. Ze positioneren het framework als een algemene oplossing die niet alleen alle varianten ondersteunt die uitdrukbaar zijn in het FlexAttention-model, maar ook meer algemene, data-afhankelijke formuleringen afhandelt die de mogelijkheden van FlexAttention overstijgen. Hoewel de huidige evaluatie zich richt op inferentie, merken de auteurs op dat de onderliggende extensies en passes algemeen doel zijn en in toekomstig werk kunnen worden uitgebreid om trainingswerklasten te ondersteunen en te optimaliseren voor bloksparsiteit.