Hugging Face ukazuje, prečo pri pozornosti v PyTorch nestačí pozerať iba na jeden riadok kódu

Hugging Face zverejnil tretí diel série o profilovaní v PyTorch a tentoraz sa zameral na pozornosť, teda attention mechanizmus, ktorý je jadrom transformerov. Text nie je produktovým oznámením nového modelu, ale má praktický význam pre tímy, ktoré ladia inferenciu alebo tréning veľkých jazykových, obrazových či video modelov. Ukazuje, že výkon attention vrstvy sa nedá bezpečne odhadnúť podľa toho, či kód vyzerá jednoducho. Rozhoduje, aké GPU jadrá sa naozaj spustia, koľko ich je a kde sa presúva pamäť.

Autori začínajú naivnou implementáciou kauzálnej pozornosti v PyTorch. V nej sú jednotlivé kroky explicitné: násobenie matíc pre dotazy a kľúče, škálovanie, maskovanie budúcich tokenov, softmax a následné násobenie hodnôt. Práve takáto verzia je didakticky užitočná, lebo v profileri vidno, ako sa každý riadok Pythonu rozpadne na CPU operácie a GPU kernely. Pre vývojára je to pripomienka, že „attention“ nie je jedna vec, ale reťazec výpočtových a pamäťových rozhodnutí.

Zaujímavý detail prichádza pri maskovaní. Bežné masked_fill vytvorí kópiu tenzora a v GPU stope sa objaví dodatočný prenos pamäte. Keď autori prejdú na in-place variant masked_fill_, tento samostatný kopírovací kernel zmizne. V malom príklade to vyzerá ako drobnosť, no v transformeri sa rovnaká operácia opakuje v mnohých vrstvách a pri dlhších sekvenciách sa podobné „drobné“ kopírovania násobia. Zároveň platí dôležitá podmienka: in-place zmena je bezpečná v ich inferenčnom príklade bez spätnej propagácie, nie automaticky v každom tréningovom scenári.

Ďalšia časť porovnáva ručne poskladanú pozornosť s torch.nn.functional.scaled_dot_product_attention. Táto funkcia vyzerá ako jedna elegantná skratka, no v skutočnosti vyberá medzi viacerými backendmi podľa dátového typu, tvaru vstupu, masky a hardvéru. Hugging Face profilu ukazuje, že matematický SDPA backend môže byť v konkrétnom nastavení výrazne pomalší než očakávanie, pretože používa viacero samostatných operácií a pomocných krokov. Jednoriadkový zápis teda nemusí znamenať optimálnu cestu.

Na opačnom konci sú fused backendy, napríklad efficient attention, flash attention alebo cuDNN attention. Tie sa snažia spojiť viac krokov do jedného jadra alebo minimalizovať pohyb dát medzi globálnou pamäťou a výpočtovými jednotkami. Praktický prínos je v tom, že väčšia časť práce zostáva v rýchlejšej pamäti a menej času sa stráca spúšťaním mnohých malých kernelov. Pre produkčné systémy to môže znamenať nižšiu latenciu odpovede, vyšší throughput alebo možnosť obslúžiť dlhší kontext na rovnakom hardvéri.

Článok je cenný aj tým, že vysvetľuje, prečo niektoré metriky v profileri môžu klamať. Napríklad nízka obsadenosť GPU nemusí automaticky znamenať zlý výkon, ak fused kernel robí presne to, čo má, a šetrí pamäťové prístupy. Rovnako veľký CPU blok v stope nemusí byť chyba modelu, ale dôsledok toho, ako konkrétny backend pripravuje alebo spúšťa generované jadro. Bez čítania celej stopy môže tím optimalizovať nesprávne miesto.

Pre slovenské a európske tímy, ktoré nasadzujú otvorené modely na vlastnej infraštruktúre, je odkaz pomerne praktický: profilovanie má byť súčasťou vývoja, nie posledná záchrana po prekročení rozpočtu za GPU. Pred nasadením novej verzie modelu alebo zmenou knižnice sa oplatí overiť, ktorý attention backend sa používa, či sa nezmenil dátový typ, či maskovanie nespúšťa zbytočné kópie a či sa očakávaná optimalizácia prejavila aj v reálnej stope.

Hugging Face zároveň neprináša univerzálny recept typu „vždy používajte jeden backend“. Skôr učí spôsob rozmýšľania: najprv si vytvoriť hypotézu, čo by mal profiler ukázať, potom stopu otvoriť a hľadať rozdiely. Pri dnešných modeloch sú práve tieto rozdiely často miestom, kde sa strácajú milisekundy, pamäť a peniaze. V ére dlhého kontextu a agentických workflow už optimalizácia attention vrstvy nie je akademický detail, ale jedna zo základných prevádzkových disciplín AI infraštruktúry.