La terza parte della serie Profiling in PyTorch di Hugging Face prende l’attention, l’algoritmo al cuore dell’architettura Transformer, e la smonta sotto il profiler. L’obiettivo non è elencare ogni trucco di ottimizzazione, ma imparare a leggere le trace: cosa lancia davvero la GPU, dove si nasconde il lavoro e perché a volte una trace più “pulita” non significa meno lavoro. Per chi lavora su inference e serving di modelli, è il tipo di competenza che separa una scelta di backend fatta a intuito da una fatta guardando i numeri. I test girano su una GPU NVIDIA A100-SXM4-80GB.
Dall’attention a mano ai primitivi
L’attention naive costruita a mano (matmul, scaling, masking causale, softmax, matmul) lancia 6 kernel per forward. La trace rivela un ospite inatteso: un Memcpy generato dall’operazione out-of-place masked_fill. Sostituendola con la versione in-place masked_fill_ il kernel di copia sparisce e si passa a 5 kernel. Sembra poca cosa, ma in un modello con molti layer il risparmio si accumula. L’articolo ricorda anche perché l’out-of-place è il default di PyTorch: autograd deve ricordare i valori del forward per il backward, quindi l’in-place è sicuro solo sotto torch.no_grad.
I quattro backend di SDPA
F.scaled_dot_product_attention nasconde molto: fa dispatch verso backend diversi e sceglie il più veloce compatibile con gli input. Il pezzo li mette a confronto uno per uno.
Il backend math è la sorpresa: lancia 20 kernel per forward ed è circa 3,7x più lento della versione naive in-place (7,239 ms contro 1,955 ms di CUDA time medio sul forward). Il motivo emerge dai nomi dei kernel: fa upcast a FP32 e gira sui CUDA core invece dei Tensor Core, ricostruisce la maschera causale a ogni chiamata e usa _safe_softmax per evitare NaN. È la reference implementation, corretta e sicura, non veloce.
I backend efficient, flash e cuDNN collassano tutto in un unico kernel fuso che resta in bf16 sui Tensor Core. L’efficient usa fmha_cutlassF (nato da xformers di Meta), il flash usa pytorch_flash (FlashAttention-2 di Tri Dao). Un punto controintuitivo: flash è il più veloce ma il profiler riporta un’occupancy bassa (circa 13%), perché ogni block usa deliberatamente molti registri e shared memory per tenere le tile dell’attention on-chip e non scrivere mai la matrice completa in HBM.
Il backend cuDNN genera un kernel su misura per il problema (simile allo spirito di torch.compile): niente transpose, lancio via cuLaunchKernelEx, ma il costo si sposta sulla CPU (circa 214 µs per forward, contro 138 di flash e 117 di efficient) per la selezione del “knob” a ogni chiamata.
In sintesi
- L’abitudine chiave della serie: indovinare prima cosa mostrerà la trace, poi guardarla e trattare ogni discrepanza come l’informazione più interessante.
- Una trace più pulita non implica meno lavoro: a volte il lavoro si sposta dentro la libreria o sulla CPU, dove il profiler lo mostra come una singola barra opaca.
- Occupancy bassa non significa kernel inefficiente: flash spende registri e shared memory apposta per evitare il traffico verso la memoria HBM.
Fonte: Profiling in PyTorch (Part 3): Attention is all you profile — https://huggingface.co/blog/torch-attention-profile
Hai qualcosa da aggiungere? Unisciti alla discussione.