Profiling en PyTorch: optimiza Attention y backends GPU

La atención es una caja negra fascinante hasta que la perfilas. En esta entrega técnica repasamos trazas reales, kernels GPU y decisiones que cambian tiempo y memoria a nivel de capa. ¿Qué pasa cuando pasas de un NaiveCausalAttention hecho a mano a F.scaled_dot_product_attention y a cada backend disponible? Vamos a verlo paso a paso.

¿Por qué perfilar attention?

Porque attention parece simple en pseudocódigo, pero en la práctica mueve mucha memoria. Matmuls, escalado, enmascarado y softmax son primitivas conocidas, pero la forma en que se ejecutan en GPU define si el modelo va rápido o se queda esperando. ¿Tu modelo sufre cuellos de botella? El profiler te responde con nombres de kernels y tiempos reales.

Naive attention y la sorpresa del memcpy

Esta es la implementación clara y directa:

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

Al perfilarlo vemos exactamente lo que esperamos: matmul, mul, masked_fill, softmax, matmul. Pero el trazo GPU revela un kernel extra: un Memcpy. ¿De dónde viene? De las operaciones out-of-place. masked_fill suele crear una copia intermedia, y esa copia aparece como memcopy en el perfil.

¿La solución? Usar la versión en sitio cuando sea seguro:

scores.masked_fill_(mask, float('-inf'))

Con ese cambio la Memcpy desaparece del trazado y ganamos tiempo y memoria. Ojo: en-place es seguro solo cuando ejecutas forward sin gradientes o sabes que no vas a corromper valores necesarios para el backward.

El one-liner SDPA que esconde varios caminos

PyTorch ofrece una llamada que lo reemplaza todo:

F.scaled_dot_product_attention(q, k, v, is_causal=True)

Genial, ¿no? Pero el profiler muestra sorpresas según el backend que use esa función. Fijémonos en el backend math: en vez de simplificar, lanza unas 20 kernels por paso y resulta aproximadamente 3.7x más lento que nuestra versión naive in-place. ¿Por qué? Porque el backend math es el de referencia: hace upcast a FP32 para seguridad numérica, reconstruye la máscara en cada llamada y usa _safe_softmax para evitar NaN en filas completamente enmascaradas. Es correcto y robusto, pero lento.

Lección: una API limpia no garantiza el camino más rápido. A veces la conveniencia traslada trabajo extra a sitios menos visibles.

Backends fusionados: efficient, flash y cuDNN

Todos buscan la misma idea: evitar escribir la matriz completa scores en la memoria global y mantener el trabajo on-chip en bf16 aprovechando Tensor Cores. Pero cada uno lo hace a su manera.

efficient (xformers)

Lanza un único kernel llamado algo como fmha_cutlassF_bf16_aligned_64x64_rf_sm80. Es un kernel fused multi-head attention construido sobre CUTLASS, opera en bf16 sin upcast y mantiene ventanas en registros y shared memory. Resultado: un solo kernel por forward y gran eficiencia.

flash (FlashAttention-2)

Vendored como pytorch_flash, implementa el truco de recorrer K y V en tiles, mantener un softmax online y nunca materializar la matriz [seq, seq]. El profiler puede mostrar baja ocupación (por ejemplo 13%) y eso asusta a primera vista, pero es intencional. Flash usa muchos registros y shared memory por bloque para mantener datos on-chip. Baja ocupación no significa ineficiente; significa que cada bloque es pesado y hace mucho trabajo con pocos accesos a HBM.

cuDNN

cuDNN genera un kernel por problema; su nombre largo refleja una versión generada y retunable. Dos puntos vitales:

• No necesita transposes metadata adicionales porque su generador consume el layout nativo.
• Lanza con cuLaunchKernelEx, así que el profiler puede reportar 0% de occupancy por una limitación de medición, no porque la GPU esté mal.

El detalle: cuDNN a menudo gana en GPU puro, pero desplaza tiempo al CPU para seleccionar y preparar la estrategia por cada llamada. En trazas eso aparece como una gran barra de CPU, a diferencia de flash o efficient.

Números rápidos del artículo original

Y entre la naive in-place y SDPA math hubo una diferencia grande en CUDA time avg: algo como 1.955 ms contra 7.239 ms para math en esa forma de referencia. La moraleja es que un backend "todo en ATen" intenta ser correcto y seguro, no necesariamente el más rápido.

Checklist práctico para perfilar attention

• Antes de mirar los trazos, adivina qué deberías ver. ¿Qué kernels esperas? Eso te ayuda a detectar anomalías.
• Si ves un Memcpy sospechoso, revisa operaciones out-of-place como masked_fill. Prueba la versión en sitio con cuidado.
• Prueba F.scaled_dot_product_attention pero cambia de backend para comparar math, efficient, flash y cudnn con torch.nn.attention.sdpa_kernel para fijarlos y perfilar cada uno.
• Observa dtype: si el backend upcastea a FP32, el coste de memoria puede duplicarse; bf16 + Tensor Cores es donde están las ganancias en A100.
• No te fíes solo de que el trazo se vea limpio. Trabajo oculto puede moverse al CPU o al interior de una librería y aparecer como una barra opaca.

Reflexión final

Profiling no es un privilegio de expertos en GPU. Es simplemente el hábito de mirar, adivinar y preguntar por las diferencias hasta que todo hace sentido. En este recorrido vimos desde un memcpy escondido por una operación out-of-place hasta kernels fusionados que aprovechan Tensor Cores y estrategias generadas por cuDNN. ¿La tarea? Abre una traza, haz tu conjetura y ve a buscar la explicación. Esa práctica te dará wins de latencia y memoria que realmente importan en modelos grandes.

Fuente original

https://huggingface.co/blog/torch-attention-profile