Hugging Face profilerer attention: math-backenden er 3,7 ganger tregere enn naiv PyTorch-kode

«Jobben er ikke å være rask, men å alltid virke.» Slik beskriver Hugging Face math-backenden i tredje og siste del av bloggserien «Profiling in PyTorch». Serien kjører attention på en NVIDIA A100 med 80 GB minne og leser profiler-tracen operasjon for operasjon. Da forfatterne byttet ut fem håndskrevne operasjoner med ett kall til F.scaled_dot_product_attention, pinnet til math-backenden, steg CUDA-tiden per forward-pass fra 1,955 til 7,239 millisekunder.

Årsaken ligger ikke i matmul-ene. Math-backenden er referanseimplementasjonen, og den er bygget for å alltid gi riktig svar. Den oppkonverterer bf16-tensorer til FP32 og havner dermed på de vanlige CUDA-kjernene i stedet for tensorkjernene. Med is_causal=True bygger den den kausale masken på nytt for hver eneste forward-pass. Den bruker _safe_softmax for å unngå NaN i rader som er fullstendig maskert, noe den håndskrevne versjonen aldri bryr seg om. Korrekt, men dyrt: 20 GPU-kjerner mot fem.

De tre fusjonerte backendene kollapser hele pipelinen til én kjerne som blir værende i bf16. Flash er raskest med 146,8 mikrosekunder CUDA-tid, foran cuDNN på 186,3 og efficient på 277,9. Flash rapporterer bare rundt 13 prosent okkupans, et tall som ser galt ut helt til du skjønner at hver blokk med vilje beslaglegger registre og delt minne for å holde attention-tiles på brikken. cuDNN viser den motsatte fellen: tracen ser renere ut uten en eneste transpose, men CPU-tiden stiger til 214 mikrosekunder fordi biblioteket genererer og tuner en ny kjerne for hvert kall.

«Når en trace plutselig blir renere, har ikke arbeidet alltid forsvunnet. Noen ganger har det bare flyttet seg dit profilereren ikke kan bryte det ned» — Hugging Face, «Profiling in PyTorch (del 3)»

Hva bør du gjøre?

1. La SDPA velge backend selv i produksjon. Math-backenden er en korrekthetsreferanse, ikke en default du vil pinne.
2. Mål CPU-tid ved siden av CUDA-tid. cuDNN-eksempelet viser at færre ATen-operasjoner i tracen kan bety mer arbeid, ikke mindre.
3. Bytt masked_fill til masked_fill_ hvis du skriver attention for hånd og kjører under torch.no_grad. Det fjerner en Memcpy-kjerne per forward-pass. Med autograd på vil in-place overskrive verdier backward-passet trenger.