Snelle attention-kernels en operatorfusie herschikken de berekening van transformers om databewegingen en kernelstarts op de accelerator te beperken. Deze optimalisaties laten in principe de gewichten van het model en zijn wiskundige functie ongewijzigd, hoewel de volgorde van bewerkingen en de precisie kleine numerieke verschillen kunnen veroorzaken. De tutorial gepubliceerd door Machine Learning Mastery gaat in detail op deze technieken in, met codevoorbeelden.
- Snelle attention-kernels en operatorfusie beperken databewegingen en kernelstarts op de accelerator.
- Fusie combineert bewerkingen in één kernel, maar helpt vooral wanneer geheugenverplaatsingen zwaar wegen.
- FlashAttention verwerkt aandacht per blok, en de gepagineerde KV-cache beperkt fragmentatie.
- torch.compile met CUDA Graphs vult fusie aan, maar vereist stabiele geheugenadressen en profilering blijft nodig.
Een transformerlaag steunt op gangbare bewerkingen: lineaire projecties, aandacht, normalisatie, activatiefuncties en residuele optellingen. Een directe implementatie start een aparte rekenkernel, een ‘kernel’, voor elke bewerking, en schrijft elke tussenliggende tensor naar het geheugen van de accelerator. De berekening blijft correct, maar de uitvoeringstijd kan meer besteed worden aan het verplaatsen van waarden en het starten van kernels dan aan nuttige berekeningen.
Operatorfusie
Het SwiGLU-blok, gebruikt door veel decoder-only modellen, illustreert dit probleem. Een directe implementatie start matrixvermenigvuldigingen voor de projecties, een kernel voor de SiLU-functie, een andere voor de elementgewijze vermenigvuldiging, en dan een laatste matrixvermenigvuldiging. Bij elke grens kan een tussenliggende tensor naar het geheugen met hoge bandbreedte (HBM) worden geschreven en daarna weer worden gelezen. Fusie bestaat uit het combineren van deze stappen in één enkele kernel. Het bespaart tijd wanneer geheugenverplaatsingen en kernelstarts zwaar wegen, en veel minder wanneer een grote matrixvermenigvuldiging de uitvoering al domineert. De keuze van de kernel blijft empirisch: de volledige werklast moet worden geprofileerd.
Andere fusies zijn gangbaar in transformers, zoals het toevoegen van een residuele verbinding gevolgd door een RMSNorm- of LayerNorm-normalisatie, of het combineren van dequantisatie met een matrixvermenigvuldiging. Fusie heeft echter zijn grenzen: bewerkingen met onverenigbare datastructuren of een zeer volumineuze tussenliggende toestand combineren slecht. Een te grote kernel kan de druk op de registers verhogen, het bezettingspercentage verlagen en trager worden dan meerdere goed afgestelde kernels.
FlashAttention-type aandacht en gepagineerde aandacht
Een klassieke aandacht berekent een scorematrix van grootte T² per hoofd, waarbij T de lengte van de sequentie is. Op een GPU kost het produceren, naar het globale geheugen schrijven en daarna weer lezen veel bandbreedte, wat een knelpunt kan worden tijdens lange ‘prefill’-fasen. FlashAttention verwerkt queries, keys en values per blok. Voor elk blok berekent het een deel van de scores, werkt het de statistieken van de softmax bij en accumuleert het de bijbehorende uitvoer in het on-chip geheugen, zonder de volledige matrix op te slaan. PyTorch stelt deze optimalisatie beschikbaar via de functie scaled_dot_product_attention(). De tutorial preciseert dat het hardcoden van een backend de code kan doen falen als geen enkele backend de invoeren of de hardware ondersteunt.
De KV-cache, die reeds berekende keys en values bewaart, groeit token voor token, voor een onvoorspelbare uiteindelijke lengte. Een aaneengesloten allocatie per query laat dus gereserveerde en ongebruikte ruimte over. De gepagineerde cache verdeelt de opslag in blokken van vaste grootte en steunt op een bloktabel die logische posities aan fysieke blokken koppelt. De aandachtskernel moet dan elke positie via deze tabel vertalen. Paging vermindert niet het aantal logische elementen, maar het beperkt fragmentatie en gereserveerde capaciteit.
In het gegeven voorbeeld gebruikt vLLM het model Qwen/Qwen2.5-1.5B-Instruct met blokken van 16 tokens (block_size=16) en de backend ‘TRITON_ATTN’. Volgens de tutorial controleert vLLM de compatibiliteit van de backend met het model en de hardware tijdens de initialisatie, en geeft het een foutmelding bij incompatibiliteit.
Fusie en CUDA Graphs met torch.compile
De functie torch.compile() legt een reeks bewerkingen vast met TorchDynamo, en compileert deze vervolgens met TorchInductor, zijn standaard backend. TorchInductor kan een enkele Triton-kernel genereren voor compatibele bewerkingen die de standaardmodus apart zou starten. Grafiekonderbrekingen, veroorzaakt bijvoorbeeld door datastroomafhankelijke controle of conversies van tensors naar Python-scalars, beperken de gebieden die de compiler kan optimaliseren.
De modus ‘reduce-overhead’ vraagt TorchInductor om CUDA Graphs te gebruiken wanneer het vastgelegde grafiek hiervoor in aanmerking komt. Een CUDA Graph registreert de reeks GPU-bewerkingen en speelt deze af met een enkele inzending aan de processor. Het creëert geen extra fusie, het vult die van TorchInductor aan. Het vereist stabiele geheugenadressen en compatibele vormen, en het kan meer geheugen verbruiken. De tutorial raadt aan de resultaten te controleren met toleranties aangepast aan het datatype, de uitvoering te profileren en de end-to-end prestaties te meten.
