Hoe vermindert Flash Attention het geheugentransport tijdens lokale inferentie?

Flash Attention vermindert het geheugentransport door exacte attention in tegels op te delen, zodat tussenliggende scores op de chip kunnen worden verwerkt in plaats van steeds naar de GPU-HBM te worden geschreven en daaruit te worden gelezen.

Op een lokale AI-server wordt dit effect belangrijker naarmate de prompt langer wordt, omdat standaard attention steeds grotere tussenliggende bewerkingen creëert. Flash Attention verandert de kennis van het model niet en verwijdert de KV-cache niet; het reorganiseert de attention-berekening om minder gegevens door de traagste niveaus van de GPU-geheugenhiërarchie te verplaatsen.

Standaard attention verplaatst grote tussenliggende tensors via HBM

Attention begint met query-, key- en value-tensors en moet scores berekenen, deze normaliseren en ze combineren met de values. Een eenvoudige implementatie materialiseert grote tussenliggende matrices in geheugen met hoge bandbreedte tussen deze stappen.

Het oorspronkelijke FlashAttention-paper stelt dat de prestaties van attention niet alleen worden beperkt door rekenwerk, maar ook door lees- en schrijfbewerkingen tussen HBM en on-chip SRAM. Lange reeksen versterken dit dataverkeer, omdat de scorematrix groeit met het aantal tokenparen.

Een thuis-GPU kan daardoor aanzienlijke tijd besteden aan het verplaatsen van gegevens in plaats van aan nuttige matrixberekeningen, vooral wanneer de contextlengte toeneemt.

Flash Attention deelt het probleem op in on-chip werksets

Flash Attention verwerkt blokken met queries, keys en values die groot genoeg zijn om in snel on-chip geheugen te passen. Elke tegel voert meerdere attention-stappen uit voordat tussenliggende gegevens de chip verlaten.

FlashAttention-3 beschrijft de kernaanpak van deze familie als een strategie met tegels die tussenliggend verkeer naar globaal geheugen voorkomt. Het algoritme produceert nog steeds exacte attention voor de ondersteunde formulering.

De tegelgrootte hangt af van de hardwarebronnen en de implementatie van de kernel. Het belangrijkste mechanisme is datalokaliteit: geladen waarden hergebruiken zolang ze zich dicht bij de rekeneenheden bevinden.

Online softmax voorkomt het opslaan van de volledige scorematrix

Bij gewone softmax lijken alle scores van een rij nodig te zijn voordat normalisatie kan plaatsvinden, waardoor de implementatie een grote matrix zou moeten materialiseren. Flash Attention houdt in plaats daarvan lopende normalisatiestatistieken bij over meerdere tegels.

De fused-attentiontutorial van Triton demonstreert gefuseerde attention met blokgewijze softmax-status. Hierdoor kunnen gedeeltelijke scoreblokken worden verwerkt en verwijderd zonder elke tussenliggende waarde terug naar HBM te schrijven.

Daar komen minder geheugentransport en minder tijdelijk geheugen samen: het algoritme herschikt de berekening zo dat een groot tussenliggend object nooit in globaal geheugen hoeft te bestaan.

Featured

ZimaCube 2 Persoonlijke Cloud Home NAS

ZimaCube2

$799.00 - $2,499.00

Buy Now

Kernel-fusie houdt meer van de attention-keten binnen één uitvoeringspad

Scoreberekening, schaling, maskering, normalisatie en value-aggregatie kunnen worden gefuseerd in één gecoördineerde kernel, in plaats van afzonderlijke kernels te starten die volledige tussenliggende resultaten via globaal geheugen uitwisselen.

PyTorchs interface voor scaled dot-product attention kan gefuseerde attention-implementaties selecteren wanneer het apparaat, de dtype, de vorm en de gevraagde functies worden ondersteund.

Het resultaat kan een snellere promptverwerking en een lager piekgebruik van tijdelijk geheugen zijn, maar de precieze winst hangt ervan af of de runtime daadwerkelijk een flash-compatibele kernel gebruikt.

Minder IO is niet hetzelfde als het elimineren van de KV-cache

Tijdens autoregressieve inferentie moeten eerder gegenereerde keys en values beschikbaar blijven voor latere tokens, tenzij het model of de runtime een afzonderlijk mechanisme voor verwijdering, compressie, venstering of offloading gebruikt.

De ROCm-richtlijnen van AMD beschrijven Flash Attention als een verbetering van de geheugenlokaliteit tussen on-chip opslag en HBM. Deze optimalisatie richt zich op de manier waarop attention actieve status leest en verwerkt, niet op de vraag of historische KV-status bestaat.

Dit onderscheid voorkomt een veelgemaakte planningsfout: Flash Attention kan werkgeheugen en IO verminderen, terwijl een server met een lange context toch zonder capaciteit komt te zitten omdat de behouden KV-cache blijft groeien.

Hardware- en vormondersteuning bepalen of het snelle pad wordt gebruikt

Flash-kernels zijn gespecialiseerd. De hoofddimensie, dtype, het attention-masker, de GPU-architectuur, de runtimeversie en andere functies kunnen een terugval naar een andere implementatie afdwingen.

Het artikel van ZimaSpace over attentiongeheugen versus modelparameters behandelt de capaciteitskant. Flash Attention moet afzonderlijk worden gemeten als optimalisatie van het uitvoeringspad.

Benchmark korte en lange prompts, controleer welke backend is geselecteerd en noteer de prefill-latentie, het maximaal toegewezen geheugen, de geheugenbandbreedte en de outputsnelheid in tokens. De hardwaregids voor lokale AI van ZimaSpace biedt de bredere capaciteitsbasis die kerneloptimalisatie niet kan vervangen.

Transformer Engine van NVIDIA bespreekt geoptimaliseerde attention-backends en hardwareafhankelijke uitvoering als onderdeel van een bredere stack voor kernelselectie.