🤖 machine learning

Hardware-Aware FP4 FlashAttention-4

Dit artikel introduceert Hardware-Aware FP4 FlashAttention-4, een methode die de beperkingen van Blackwell's FP4 tensor cores overwint door Direct-P te gebruiken voor niet-causale inferentie en een causaal pad met FP8 gradiënten, waarmee tot 2,13× snellere forward throughput en 1,14× snellere single-GPU training updates worden bereikt terwijl de divergentieproblemen die bij MXFP4 training worden waargenomen, worden vermeden.

Oorspronkelijke auteurs: Robert Hu

Gepubliceerd 2026-09-04✓ Author reviewed ⓘ

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Robert Hu

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 door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

In de wereld van kunstmatige intelligentie vertrouwen de krachtigste modellen op een mechanisme genaamd "attention" (aandacht) om context te begrijpen. Stel je een lezer voor die een lang document scant; aandacht stelt het model in staat om zich te concentreren op de meest relevante woorden in een zin terwijl de rest wordt genegeerd, waarbij ideeën over grote afstanden met elkaar worden verbonden. Om dit te doen, voert de computer een enorme reeks berekeningen uit, waarbij elk woord met elk ander woord wordt vergeleken om hun relaties te bepalen. Jarenlang zijn deze berekeningen uitgevoerd met hoogprecisiegetallen, vergelijkbaar met het gebruik van een liniaal met piepkleine, precieze markeringen om een gebouw te meten. Dit zorgt ervoor dat het model correct leert, maar het is traag en verbruikt enorme hoeveelheden energie. Naarmate modellen groter worden, wordt de behoefte aan snelheid cruciaal. Onderzoekers hebben recentelijk chips ontwikkeld die in staat zijn om veel kleinere, grovere getallen te gebruiken—vierbits floating-point-waarden—om deze vergelijkingen veel sneller uit te voeren. Het simpelweg overstappen naar deze kleinere getallen is echter niet genoeg; de software die de datastroom beheert, moet ook veranderen, anders verdwijnen de snelheidswinsten.

Een onderzoeker bij Graphcore heeft dit specifieke knelpunt aangepakt met een nieuwe methode die zij Direct-P noemen. Hun werk richt zich op de "forward" pass (voorwaartse stap) van attention, waarbij het model informatie verwerkt om een antwoord te genereren, en de "backward" pass (achterwaartse stap), waarbij het leert van zijn fouten. De onderzoeker ontdekte dat hoewel de nieuwe chips getallen ongelooflijk snel kunnen vermenigvuldigen, de stap daartussenin—het omzetten van ruwe scores naar waarschijnlijkheden—alles vertraagde. Het was alsof er een super-snelle lopende band was die steeds stopte omdat een arbeider elk onderdeel zorgvuldig moest opmeten voordat hij het kon doorgeven. De onderzoeker ontdekte dat de standaardmanier om deze conversie af te handelen, die complexe schaling en afronding inhield, een verkeersopstopping veroorzaakte die de snelheidsvoordelen van de nieuwe hardware tenietdeed.

Om dit op te lossen, heeft de onderzoeker het conversieproces herontworpen om direct en gestroomlijnd te zijn. In plaats van een precieze waarschijnlijkheid te berekenen en deze vervolgens af te ronden, brengt hun methode de ruwe scores direct in kaart naar de specifieke codes die de hardware gebruikt. Dit elimineert de tussenstappen die de vertragingen veroorzaakten. Toen ze deze aanpak testten op de nieuwste generatie data-center chips, waren de resultaten opmerkelijk. Voor bepaalde veelvoorkomende vormen van data verwerkte de nieuwe methode informatie meer dan twee keer zo snel als de vorige standaard, die vertrouwde op getallen met een hogere precisie. Dit werd bereikt terwijl de output nauwkeurig genoeg bleef voor complexe taken zoals videogeneratie en taalbegrip. In tests met een videogeneratiemodel was de nieuwe methode bijna twee keer zo snel als de standaardaanpak, waarbij de resultaten, hoewel eindig en bruikbaar, een meetbare drift en lagere gelijkenisscores vertoonden vergeleken met de tragere, meer precieze versie.

Het verhaal gaat echter niet alleen over snelheid; het gaat ook over wat er gebeurt wanneer het model probeert te leren. De onderzoeker onderzocht of zij deze ultra-snelle, lage-precisiegetallen voor het gehele trainingsproces konden gebruiken, inclusief de backward pass waar het model zijn kennis bijwerkt. Zij ontdekten dat hoewel de forward pass op volle snelheid kon draaien, de backward pass een zorgvuldige afweging vereiste. Wanneer zij probeerden de snelste vierbits-indeling voor de waarschijnlijkheidswaarden te gebruiken tijdens de training, werd het leerproces onstabiel en slaagde het model er niet in om te verbeteren. De getallen werden te grof, waardoor het leersignaal uiteenviel. Bijgevolg bepaalde de onderzoeker dat voor een stabiele training een iets preciezere achtbits-indeling voor de waarschijnlijkheidswaarden moet worden gebruikt, zelfs als de rest van de berekening de snellere vierbits-indeling gebruikt. Deze hybride aanpak stelde hen in staat om de volledige trainingscyclus met ongeveer 14 procent te versnellen op een enkele krachtige chip, een aanzienlijke winst voor grootschalige modellen.

De onderzoeker keek ook naar de fysieke limieten van de computerchips zelf. Zij ontdekten dat de snelheid van de berekening niet langer het hoofdprobleem was; het probleem was hoe de data werd opgeslagen en rondbewogen binnen de chip. De chip heeft een klein, ultra-snel geheugengebied waar de getallen worden vastgehouden terwijl er aan gewerkt wordt. De onderzoeker vond dat deze geheugenruimte volledig bezet was, waardoor er geen ruimte was om verschillende stadia van de berekening te overlappen. Het was als een keuken waar het aanrecht zo vol staat met ingrediënten dat de chef niet kan beginnen met het snijden van de volgende groente totdat de huidige klaar is, zelfs als het mes ongelooflijk scherp is. Omdat de geheugenruimte volledig bezet was, kon de chip niet aan meerdere stappen tegelijk werken, wat beperkte hoeveel sneller het hele proces kon worden.

Dit werk benadrukt een cruciale verschuiving in hoe we denken over het sneller maken van kunstmatige intelligentie. Het is niet genoeg om simpelweg snellere rekenmachines te bouwen; de volledige workflow moet worden herontworpen om aan te sluiten bij de mogelijkheden van de hardware. De onderzoeker liet zien dat door de manier waarop waarschijnlijkheden worden berekend te veranderen en door de datastroom zorgvuldig te beheren, zij enorme snelheidsverbeteringen konden ontsluiten. Toch bewezen zij ook dat er harde grenzen zijn. De geheugenbeperkingen van de chip betekenen dat zelfs met de snelst mogelijke wiskunde, er een plafond is aan hoeveel overlap mogelijk is. De weg vooruit, suggereren zij, ligt niet alleen in snellere aritmetiek, maar in slimmere manieren om de data te organiseren zodat de middelen van de chip nooit onbenut blijven wachten. Deze balans tussen brute snelheid en zorgvuldig databeheer is wat de volgende generatie kunstmatige intelligentie efficiënt zal laten draaien.

=== SAMENVATTING ===

Technische Samenvatting: Hardware-bewuste FP4 FlashAttention-4

Probleemstelling

NVIDIA Blackwell GPU's introduceren 4-bit floating-point (FP4) tensor cores die matrixvermenigvuldigingen aanzienlijk versnellen. Echter, in attention-mechanismen wordt de versnelling vaak niet beperkt door de matrixproducten ( $QK^T$ en $PV$ ), maar door de tussenliggende softmax-fase. Deze fase omvat score-reductie, exponentiële evaluatie, waarschijnlijkheidsconstructie en schaling, wat voltooid moet zijn voordat het tweede matrixproduct kan beginnen.

Bestaande pogingen om FP4 voor de volledige attention-blok te benutten, stuiten op twee primaire hindernissen:

Timing: Het construeren van de waarschijnlijkheidstensor ( $P$ ) uit scores ( $S$ ) creëert een kritiek pad-afhankelijkheid die de $PV$ -product vertraagt.
Numeriek Bereik: FP4-formaten (specifiek NVFP4 en MXFP4) gebruiken een 4-bit payload (E2M1) met een beperkt dynamisch bereik. Standaard softmax-waarschijnlijkheden beslaan een bereik dat onhandig is voor deze payload, wat leidt tot onderflow, "zero scales" (waarbij volledige blokken verdwijnen) of aanzienlijk precisieverlies als dit niet zorgvuldig wordt beheerd.

Verder, in causale training, vereist de backward pass doorgaans het reconstrueren van waarschijnlijkheden uit opgeslagen staten. Benaderingen die volledige precisie-scores opslaan of deze opnieuw berekenen in BF16, ontzeggen het potentieel voor versnelling bij training met lage precisie.

Methodologie: Direct-P en Causal Backward

Het artikel stelt Direct-P voor, een methode om volledige FP4 forward inference mogelijk te maken en een gekwantiseerde causale backward path op Blackwell-hardware (GB200/B300).

1. Direct-P voor Forward Inference

Direct-P modificeert het interval tussen de voltooiing van de $QK^T$ score-fragment en de uitgave van de $PV$ instructie. Het behoudt de bestaande HAO AI Lab two-query pipeline en Tensor Memory (TMEM) ownership protocol, maar optimaliseert de waarschijnlijkheidsgeneratie:

Directe Mapping: In plaats van een high-precision exponent te berekenen en vervolgens af te ronden, behandelt Direct-P de constructie van waarschijnlijkheid als een classificatieprobleem. Het brengt genormaliseerde log-scores direct in kaart naar de 8 mogelijke niet-negatieve E2M1-magnitudes met behulp van een affine classifier (of selectieve native exponenten). Dit verkort het kritieke pad.
Consistente Normalisatie: De noemer (normalisator) voor de softmax wordt direct geaccumuleerd vanuit de afgeronde FP4-waarschijnlijkheidscodes en block scales die door de $PV$ instructie zullen worden geconsumeerd. Dit zorgt ervoor dat de teller en de noemer dezelfde benaderde operator beschrijven, waardoor mismatches tussen high-precision sommen en low-precision payloads worden voorkomen.
Range Guarding: Voor lagen met extreme logits (bijv. >500) voorkomt een sampled anchor-mechanisme subnormale underflows zonder dat een volledige re-scan van de scores nodig is.
Formaatselectie: De methode gebruikt NVFP4 voor $Q$ en $K$ (vanwege fijne schaling) en MXFP4 voor $P$ en $V$ (vanwege de power-of-two block scales die goed aansluiten bij de $N_{32}$ producer fragments).

2. Gekwantiseerde Causal Backward

Voor training introduceert het artikel een pad waarbij de forward pass de gekwantiseerde $Q/K$ payload, de schalen en de log-sum-exp (LSE) normalisator opslaat.

Reconstructie: De backward pass reconstrueert waarschijnlijkheden direct uit deze opgeslagen gekwantiseerde operanden in plaats van een apart BF16 score-pad te creëren.
Gradiëntformaten: De projectie-epilogen publiceren rij/kolom-georiënteerde FP8 views van $Q, K, V$ om de gradiënt-matrixproducten te voeden.
Stabiliteitsbeperking: Hoewel de forward path MXFP4 gebruikt voor $P/V$ , toonden experimenten met gedistribueerde training aan dat MXFP4 P/V trajecten divergeren. Daarom gebruikt de behouden trainingsroute FP8 voor $P$ en $V$ in de backward pass om stabiliteit te garanderen, terwijl $Q/K$ in NVFP4 blijft.

Belangrijkste Bijdragen

Direct-P Implementatie: Een methode die scores direct in kaart brengt naar MXFP4 waarschijnlijkheidscodes en normaliseert met behulp van dezelfde gerepresenteerde waarden, wat een doorvoer van tot wel 2.13× BF16 forward op GB200 bereikt voor gunstige vormen, met een geometrisch gemiddelde versnelling van 2.02× over de geteste D128 grid.
Gekwantiseerde Causal Training Path: Een backward pass die waarschijnlijkheden reconstrueert uit opgeslagen forward kwantisatiestaten, wat een volledige 8-miljard parameter single-GPU update met maximaal 1.14× versnelt.
Hardware Profiling: Identificatie van Tensor Memory (TMEM) ownership als de primaire limiet op overlap. De studie toont aan dat het louter hebben van snellere tensor cores onvoldoende is als de opslaglay-out (512 kolommen bij D128) voorkomt dat de volgende $QK$ operatie overlapt met de huidige $PV$ fase.

Resultaten

Forward Inference Prestaties

Doorvoer: Op GB200 bereikt de "snelle" Direct-P route een geometrisch gemiddelde versnelling van 2.02× ten opzichte van de HAO BF16 baseline, met pieken van 2998 TFLOP/s (S32768/H24).
Nauwkeurigheid: Direct-P ruilt enige precisie in voor snelheid. Op Gaussian softmax tests heeft de "snelle" route een gemiddelde cosine similarity van ~0.944 tegenover BF16, vergeleken met ~0.990 voor de HAO NV/FP8 route.
Modelimpact:
- ViT/BERT: Bij evaluaties met vaste input behoudt de "snelle" route de taaknauwkeurigheid (bijv. ViT S4096 top-1 nauwkeurigheid komt overeen met BF16 op 88.5%) ondanks lagere operator-niveau cosine similarity.
- Wan Video Diffusion: De methode blijft finit en functioneel voor 20-staps diffusie. Echter, fouten accumuleren over stappen; het 14B model bij 20 stappen vertoont een relatieve-L2 van ~0.53 vergeleken met BF16, wat wijst op drift maar functionele generatie.

Training Prestaties

Single-GPU: De gekwantiseerde route (NVFP4 projecties + FP8 P/V) versnelt de volledige attention sublayer met 1.25× en de volledige 8B model update met 1.14× bij een batch size van 4.
Gedistribueerde Training: In een gematchte 64-GPU experiment over 100 miljard tokens bereikte de FP8 P/V route een 1.112× doorvoerverbetering ten opzichte van BF16.
Stabiliteit: Het artikel merkt expliciet op dat MXFP4 P/V trajecten divergeerden in gedistribueerde settings. De succesvolle trainingsroute vertrouwt op FP8 voor de $P/V$ product, wat betekent dat de trainingsroute niet "zuiver" FP4 end-to-end is.

Betekenis en Claims

Het artikel claimt dat Direct-P erin slaagt om volledige FP4 forward attention praktisch te maken door het kritieke pad tussen matrixproducten te verkorten. Het demonstreert dat hoewel FP4 matrix cores sneller zijn, de bottleneck is verschoven naar datamovement, waarschijnlijkheidsconstructie en memory ownership.

Belangrijkste leerpunten:

Hardware Limieten: De primaire barrière voor verdere versnelling is niet rekenkundige doorvoer, maar TMEM ownership. De huidige D128 lay-out consumeert alle 512 TMEM kolommen, wat voorkomt dat de volgende $QK$ operatie overlapt met de staart van de $PV$ fase.
Readiness Stalls: Lage tensor-core activiteit is vaak een symptoom van operand readiness (wachten op waarschijnlijkheidsconstructie) in plaats van een gebrek aan rekenkracht.
Training Trade-offs: Hoewel FP4 forward inference levensvatbaar is, vereist stabiele causale training momenteel een concessie aan FP8 voor de waarschijnlijkheids/waarde operanden om divergentie te voorkomen. Het artikel suggereert dat toekomstige formaten zoals Unsigned E5M3 (UE5M3) een volledige FP4 training mogelijk kunnen maken, maar dit wordt hier niet geëvalueerd.

De conclusie van het werk is dat snellere matrix arithmetic alleen onvoldoende is; het bereiken van hogere efficiëntie vereist architecturale veranderingen die grotere allocabele overlap windows (bijv. extra score bestemmingen) mogelijk maken naast geoptimaliseerde low-precision algoritmen.