Ez a cikk a "Profiling in PyTorch" sorozat harmadik része, amelynek célja, hogy megtanuljuk olvasni a profiler nyomvonalakat és táblázatokat, és ebből kiindulva optimalizálási döntéseket hozzunk. Az előzmények: az első részben egyszerű műveleteket (összeadás, szorzás) profilzottunk, a második részben lineáris rétegeket és egy fuzionált multilayer perceptront vizsgáltunk. Itt az attention (összpontosítás) műveletét vizsgáljuk meg, mert a Transformer‑szerű modellek alapvető építőköve — ugyan kvadratikus komplexitású, de több trükk létezik a hatékonnyá tételére.
A cikkben használt script fájlok a következők: 04_a_naive_attention.py, 04_b_inplace_ops_attention.py, 04_c_sdpa_attention.py és 04_d_kernels_attention.py. A méréseket egy NVIDIA A100‑SXM4‑80GB GPU-n futtattuk. A példákban a PyTorch profiler kimenetét nézzük meg CPU és GPU sávokra bontva, és összehasonlítjuk, hogyan jelennek meg az egyes megoldások a trace‑en.
Naiv causal attention
Az attention lépései röviden: számoljuk a pontszámokat (scores = matmul(q, k.T)), skálázzuk, alkalmazunk causal maszkot (masked_fill), normalizálunk softmax‑szal, majd a súlyokkal átlagolunk (matmul(attn, v)). Ez igazából primitív műveletek sorozata — így írtunk egy egyszerű PyTorch modult:
class NaiveCausalAttention(nn.Module): def init(self, head_dim): super().init() self.scale = 1.0 / math.sqrt(head_dim) def forward(self, q, k, v, mask): scores = torch.matmul(q, k.transpose(-2, -1)) scores = scores * self.scale scores = scores.masked_fill(mask, float("-inf")) attn = torch.softmax(scores, dim=-1) out = torch.matmul(attn, v) return out
Előre megtippeltük a trace‑en láthatókat: két matmul, egy mul (skálázás), egy masked_fill, egy softmax. A CPU sávon ezt is látjuk. A GPU sávra kitekintve viszont öt várt kernel mellett egy váratlan Memcpy kernel is megjelenik — ez a masked_fill out‑of‑place viselkedéséből származik: PyTorch gyakran készít másolatot, végrehajt rajta egy műveletet és a másolatot adja vissza.
In‑place masking
Ha a masked_fill helyett a PyTorch in‑place konvenciója szerinti masked_fill_‑t használjuk (egy sor módosítása), a Memcpy kernel eltűnik a GPU trace‑ről. Tehát egyszerűen:
scores.masked_fill_(mask, float("-inf"))
eltávolít egy GPU memcpy‑t minden forward hívásból. Bár ez apró megtakarításnak tűnik, egy nagy modellnél (sok réteg, sok attention hívás) összegződve jelentős hatása lehet. Meg kell azonban jegyezni, hogy in‑place műveletek általában veszélyesek a gradiensek számítása miatt — itt a forward torch.no_grad alatt fut, így biztonságos az in‑place használat.
Scaled Dot Product Attention (SDPA) és backendes választás
A PyTorch egy beépített API‑t ad: F.scaled_dot_product_attention(q, k, v, is_causal=True). Ez a hívás több belső implementációhoz (backendhez) képes átirányítani a hívást a leggyorsabb, a bemenetet támogató megoldásra (dtypes, head_dim, maszk, hardware függvényében). A backends felsorolása például: MATH, FLASH_ATTENTION, EFFICIENT_ATTENTION, CUDNN_ATTENTION.
A tesztekhez rögzítjük a backendet és megnézzük mindegyik trace‑ét külön.
Math backend (referencia)
A math backend a referencia: primitív ATen hívásokra bontja az attentiont, számításbiztonságos és NaN‑védett. Eredmények összefoglalva:
- A profiler táblázat szerint a Naive in‑place forward átlagos CUDA ideje ~1.955 ms, míg az SDPA math backend ugyanazon opra ~7.239 ms; a teljes CUDA idő 7.194 ms vs 27.279 ms (a math backend körülbelül 3.7× lassabb).
- A traceben a math backend egy forwardra körülbelül 20 GPU kernelt indít, szemben a kézzel írt naiv attention 5 kernelével.
Miért lassú? Több ok együttállása:
- A math backend biztonság kedvéért FP32‑re upcastolja a bemeneteket (például ha bemenetek bf16), ezért több adatmozgatás történik és a matmulok a lassabb, klasszikus sgemm (CUDA magok) útvonalra mennek, nem a Tensor Core bf16 gyorsútjára.
- Ha is_causal=True, a backend újra és újra létrehozza a maszkot minden hívásnál (CPU oldalon és GPU‑n triu/tril, where, add_ műveletek), tehát a maszképítés nem tűnik el, csak áthelyeződik.
- A math backend safe_softmax‑ot használ, amely extra lépésekkel védi a NaN‑októl, további kernelindításokkal jár.
Összefoglalva: a math backend a helyes, robosztus referencia, de sok köztes eredményt ír ki és lassú.
Efficient backend (xformers alapú)
Az efficient backend egyetlen, fuzionált kernelre redukálja az egészet: fmha_cutlassF_bf16_aligned_64x64_rf_sm80. Jellemzői:
- Egy fused multi‑head attention kernel (fmha), CUTLASS alapú, fut bf16 formátumban (nincs FP32 upcast), a regiszterfájlban dolgozik, kis on‑chip mozgással.
- Ez a megoldás lényegében egy, a Tensor Core‑okat kihasználó, memóriahatékony implementáció, amely elkerüli a nagy, [seq, seq] mátrix teljes kiírását a GPU HBM‑re.
Flash backend (FlashAttention‑2)
A flash backend szintén egyetlen fuzionált kernelrel dolgozik (pytorch_flash), ez Tri Dao FlashAttention‑2 implementációja. Ez a megoldás az online softmax és a tiled feldolgozás segítségével soha nem írja ki a teljes [seq, seq] score mátrixot a GPU globális memóriájába; a számítást részblokkonként tartja on‑chip.
Érdekes jelenség: a flash kernel profilerben alacsony bejelentett occupancy‑t mutat (~13%). Ez megtévesztő lehet, de magyarázata az, hogy a kernel blokkonként sok regisztert és megosztott memóriát használ; emiatt egy SM‑en csak kevés blokk fér el egyszerre, azaz alacsony a resident warpok aránya. Ez viszont nem rossz jel: a terv az, hogy sok gyors on‑chip erőforrást fogyasztva minimalizáljuk a HBM‑hez való ki/be járást, és ez éppen a gyors működés oka.
cuDNN backend
A cuDNN backend szintén egy fuzionált kernelhívást ad vissza, de fontos különbségek vannak:
- A cuDNN kernel generált, per‑problem (shape‑függően) optimalizált kódot hoz létre, ezért nincs szükség transzponálásokra a CPU oldalon: a generator közvetlenül a [B, H, S, D] layoutot fogyasztja, míg a flash/efficient megoldások gyakran metadata transposokat hajtanak végre.
- A kernel indítása driver‑szinten történik (cuLaunchKernelEx), ami miatt a CUPTI profiler nem tudja helyesen mérni az occupancy‑t (0%‑ként jelentkezhet), de a kernel footprintje alapján az erőforrásigény hasonló a flash‑hez (például nagy regiszterhasználat, egy blokk fér el SM‑enként).
- A cuDNN hátránya itt, hogy több CPU időt használ: az mérések szerint a CPU átlagidő per forward ~214 µs, míg flash ~138 µs és efficient ~117 µs. A magyarázat az, hogy cuDNN minden híváskor kiválaszt és előkészít egy tervet/konfigurációt (knob), ez a runtime költség.
GPU oldalon a mért átlagidők ezen a mintán: efficient ~277.9 µs, flash ~146.8 µs, cuDNN ~186.3 µs (itt a flash a leggyorsabb erre a konkrét shape‑re, de különböző alakoknál a cuDNN gyakran felülmúlja a többit, mert per‑shape tuningot végez).
Mit tanultunk a trace‑ekből (összefoglaló)
- Egy kis, egyszerű kódmódosítás (masked_fill → masked_fill_) eltávolíthat egy Memcpy GPU kernel‑t per forward, és ezzel időt és memóriát takaríthat meg.
- A F.scaled_dot_product_attention gazdag backend‑választó réteget rejt: a math backend a helyes, robosztus, dtype/NaN‑biztos referencia, de sok kernelindítással és FP32 upcasttal lassú (például ~3.7× lassabb az adott példán).
- Efficient (xformers) és FlashAttention‑2 egyetlen fuzionált Tensor Core‑alapú kernelre tömörítik az attentiont és megőrzik bf16 pontosságot, így jelentősen gyorsabbak.
- A flash kernel alacsony bejelentett occupancy‑t mutathat, de ez nem feltétlenül rossz; sok regisztert és shared memory‑t használ szándékosan, hogy elkerülje a HBM‑t.
- A cuDNN backend egy generált, per‑probléma optimalizált kernelhez vezet, ami eltünteti a transzponálásokat, de több CPU munkát (pl. knob kiválasztás) helyez a hívás előkészítésére.
Következtetések
A profilozás legfontosabb gyakorlatát foglalja össze a cikk: először meg kell tippelni, mit látunk a trace‑en, aztán megnyitni a trace‑et és a tipp és a valóság közti eltérést követni. A trace‑ből származó meglepetések (a rejtett memcpy, a math backend 20 kernelje, a flash alacsony occupancy‑je, a cuDNN nagy CPU‑ideje) mind abból adódtak, hogy kiléptünk az elméleti elképzelésekből és megnéztük a tényleges futásmódot.
Profilozás nem csak GPU‑szakértőknek való titok: fegyelem, figyelem és kérdezés kérdése — "várj, miért történik ez?" — addig kérdezni, amíg az ok meg nem világosodik. Nyiss meg egy trace‑et, tégy egy tippet, és találj rá a mismatch‑re.
Köszönet Noe Flandre‑nek a korai vázlatok átnézéséért.
(A poszt szerkesztésénél LLM‑et is használtak stilisztikai javításokra; ez nem jelenti, hogy az üzenet generálását kizárólag automatára bízták.)
