Perfilado de atención en PyTorch: qué mirar y cómo optimizar

Introducción: por qué perfilar atención

Esta entrega continúa la serie sobre perfilado en PyTorch y se centra en la operación de atención, un bloque fundamental de los transformer. Hasta ahora hemos visto cómo descubrir hotspots con tablas de profiler y cómo interpretar trazas que muestran el orden y la duración de operaciones. La atención, famosa por su costo cuadrático en la longitud de secuencia, combina varias operaciones simples —matmuls, escalados, máscaras y softmax— que juntas definen su performance. Nuestro objetivo aquí es mostrar cómo esas piezas aparecen en un perfil y cómo un cambio pequeño en el código puede eliminar trabajo redundante en GPU.

Los scripts usados en este experimento están disponibles en los archivos: 04_a_naive_attention.py, 04_b_inplace_ops_attention.py, 04_c_sdpa_attention.py y 04_d_kernels_attention.py. Los resultados se registraron con una NVIDIA A100-SXM4-80GB; si desean reproducir las pruebas, Hugging Face facilita el acceso a GPU mediante Dev Mode en Spaces o la pipeline de Jobs.

Descomponiendo la atención: operaciones elementales

La atención causal clásica trabaja con tres tensores: Queries (q), Keys (k) y Values (v). El flujo típico es:

• calcular scores con una multiplicación matricial: matmul(q, k.T)
• escalar los scores: scores * scale
• aplicar una máscara causal: scores.masked_fill(mask, -inf)
• normalizar con softmax: attn = softmax(scores)
• reponder los valores: out = matmul(attn, v)

En esencia, la atención es una cadena de operaciones primitivas. Al perfilarla, esperamos ver exactamente esos pasos como kernels lanzados o llamadas CPU que disparan kernels: dos matmuls, una multiplicación escalar, la aplicación de la máscara y el softmax.

Implementación ingenua y trazas esperadas

Una versión directa en PyTorch se parece a esto (fragmento tomado de los experimentos):

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 perfilar la ejecución, en la pista de CPU vemos las llamadas que corresponden a cada operación y, al desplegar la pista de GPU, leemos los kernels activados en cada paso.

Hallazgo inesperado: una copia de memoria en GPU

Al profundizar en la pista GPU para una pasada de forward, aparecen estos kernels en orden: matmul (q·kᵀ), mul (escalado), una memcpy, causal masking, softmax y otro matmul (attn·v). Las cinco primeras operaciones eran esperadas; la copia de memoria (memcpy) no lo era.

¿Por qué aparece esa copia? La pista sugiere que la llamada a scores.masked_fill, que por defecto devuelve un tensor nuevo (operación out-of-place), induce una copia de datos que se refleja en un memcpy en GPU. En términos de profiler, eso significa que antes de aplicar la máscara PyTorch materializa una nueva región de memoria en la GPU, lo que añade latencia y uso de memcpy por cada paso de atención.

Solución simple: usar operaciones in-place

PyTorch ofrece el sufijo _ para las versiones in-place de las operaciones. Reemplazar scores = scores.masked_fill(mask, -inf) por scores.masked_fill_(mask, -inf) evita crear una copia adicional. Con ese único cambio se observa que la memcpy desaparece de la pista GPU: el trabajo de la GPU se reduce en un kernel por paso de atención.

En los experimentos se compararon las trazas del caso out-of-place y la versión in-place. En la pista CPU la sección de masking quedó mucho más ligera cuando se usó la operación in-place, y en la pista GPU la Memcpy kernel ya no aparece. Esto confirma que una pequeña modificación puede eliminar kernels redundantes que se repiten muchas veces en modelos con decenas de capas.

Por qué importa para modelos reales

En un modelo transformer grande, la operación de atención se ejecuta una vez por cabeza y por capa. Aunque ahorrar un solo kernel por paso pueda parecer marginal en una sola llamada, el efecto se multiplica a lo largo de millones de ejecuciones durante entrenamiento o inferencia. En la práctica, reducir kernels innecesarios ayuda a disminuir latencia, congestión de memoria y, potencialmente, el coste operativo en infraestructuras GPU.

Además, optimizaciones de este tipo son acumulativas: junto con técnicas como kernels fusionados, atención causal optimizada (SDPA), o kernels escritos a medida para ciertos tamaños de cabeza, cada ahorro suma para llegar a implementaciones que escalan mejor en producción.

Qué revisar cuando perfilar atención en su entorno

• Inspeccionen ambas pistas del profiler: CPU y GPU. La pista CPU muestra qué llamadas lanzan kernels; la GPU revela los kernels reales y copias.
• Busquen kernels de memcpy o transferencias recurrentes. Si aparecen, identifiquen la operación out-of-place asociada en el código.
• Prueben la versión in-place con cuidado. Las operaciones in-place reducen copias, pero pueden interactuar con el grafo de autograd; validen que el backward funcione y que no haya conflictos con otros usos del tensor.
• Reproduzcan pruebas con las formas y tamaños representativos de su carga de trabajo (batch, secuencia, tamaño de cabeza). El perfil depende mucho del shape.
• Combinen este cambio con otras optimizaciones (kernel fusion, SDPA, bibliotecas especializadas) para obtener beneficios mayores.

Cómo reproducir los trazados (resumen práctico)

En los experimentos originales se ejecutaron los scripts mencionados y se generaron trazas con herramientas de Hugging Face. Un ejemplo de comandos usados fue:

uv run 04_a_naive_attention.py uvx trace-util -f traces/ -b <hf_uname>/traces

Si usan la infraestructura de Hugging Face, pueden aprovechar Dev Mode en Spaces para acceder rápidamente a una GPU y revisar las trazas; también es posible integrar la ejecución en Jobs para procesos más reproducibles.

Conclusión

Perfilar atención en PyTorch revela que operaciones aparentemente triviales —como una llamada a masked_fill que devuelve un tensor nuevo— pueden traducirse en copias de memoria en GPU. Cambiar a la versión in-place elimina ese memcpy y reduce el número de kernels por paso de atención. Este tipo de microoptimización es especialmente relevante en modelos con muchas capas y puede mejorar la eficiencia operativa cuando se escala a producción.

Para equipos en América Latina que despliegan modelos basados en transformer, revisar estos detalles en el profiler es una forma práctica de optimizar costos y rendimiento antes de invertir en hardware adicional. Los scripts de referencia y las trazas permiten explorar estos comportamientos de forma reproducible y aplicar las correcciones en implementaciones reales.