La regla es sencilla: lo que no puedes perfilar no puedes optimizar. ¿Quieres exprimir más tokens por segundo de un LLM, quitar milisegundos a la inferencia o entender por qué tu loop de entrenamiento va más lento de lo que promete la hoja de especificaciones? En algún punto tienes que abrir un rastro del profiler y leerlo.
Pero abrir ese trazo suele sentirse como entrar a una pared de rectángulos de colores con nombres intimidantes. Esta guía técnica, primera de una serie, baja esa rampa de entrada: partimos de lo más simple, una multiplicación de matrices seguida de una suma de bias, y aprendemos a leer lo que torch.profiler nos devuelve.
Preparando el experimento y qué exporta el profiler
La función que usamos es mínima y pedagógica:
def fn(x, w, b): return torch.add(torch.matmul(x, w), b)
Anotar la región con torch.profiler.record_function('matmul_add') es opcional pero útil: facilita saltar a esa etiqueta en las trazas. El flujo típico de profiling en el script es:
- definir la región a medir (aquí
fn) y envolverla enrecord_function. - ejecutar bajo
torch.profiler.profile(...)con actividades CPU y CUDA. - llamar
prof.key_averages().table(...)para la tabla estadística. - exportar la traza con
prof.export_chrome_trace(trace_path)para verla en Perfetto.
El profiler produce dos artefactos complementarios:
- La tabla del profiler: resumen estadístico que responde "qué se lleva la mayor parte del tiempo". Ideal para encontrar hotspots.
- La traza temporal (.json): vista tiempo-espacio de CPU vs GPU, con los kernels, huecos, y dependencias.
Ejemplo de comandos usados en la guía:
uv run 01_matmul_add.py --size 64
y para runs más grandes o con warmup/compile:
uv run 01_matmul_add.py --size 4096 --warmup --compile
Leer la tabla: Self vs Total y el síntoma de overhead
Al abrir la tabla verás columnas como Self CPU time total y Self CUDA time total. Atención a esta señal frecuente:
- Si
Self CPU time totalestá en milisegundos ySelf CUDA time totalen microsegundos, el patrón típico es overhead-bound: la CPU hace más trabajo preparatorio que lo que el GPU tarda en calcular.
En el ejemplo con tamaño 64 observamos algo así como:
- Self CPU time total: 2.314 ms
- Self CUDA time total: 23.104 us
Eso significa que el GPU está casi siempre inactivo: el matmul pequeño se ejecuta rápido en GPU, pero la mayor parte del tiempo se gasta en lanzar kernels, asignar workspace o ejecutar heurísticas de cuBLAS. Solución simple: aumentar la carga por llamada (matrices más grandes o batching) o fusionar operaciones.
Con --size 4096 el balance cambia:
- Self CPU time total: 4.908 ms
- Self CUDA time total: 4.495 ms
Ambos ahora en ms, estamos en una región compute-bound donde el GPU domina: el objetivo usual cuando escalar a modelos reales.
Leer la traza: CPU lane, GPU lane y las ventanas muertas
La traza muestra dos lanes principales: CPU y GPU. Los elementos clave que debes buscar y qué significan:
-
Un espacio entre la entrada de
record_function('...')y el primeraten::*es una "ventana muerta". Puede ser allocation de workspace, heurísticas de cuBLAS o lazy loading. ¿Qué hacer? Warmup. -
El scheduler de profiling controla qué pasos se registran. Ejemplo de schedule común:
torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
wait salta inicializaciones; warmup ejecuta sin registrar; active es lo que ves en la traza.
-
Si ves un offset de varios ms entre el momento en que la CPU envía trabajo y cuando el GPU empieza, puede aparecer una entrada llamada "Activity Buffer Request" en la lane de GPU: el propio profiler refiriéndose a su buffer de eventos o a alguna solicitud de memoria.
-
Si un paso es mucho más ancho que los demás (primer ProfileStep, por ejemplo) es cold-start: agrega warmup antes de medir.
Señales en la traza y llamadas CUDA que importan
Busca estas llamadas y ten en mente su significado:
-
cudaOccupancyMaxActiveBlocksPerMultiprocessorantes de uncudaLaunchKernel: es una consulta de planificación. Suele aparecer para kernels pesados como GEMM o convoluciones. Indica que cuBLAS está eligiendo heurísticamente una variante. -
cudaLaunchKernelsin la ocupancy query: típicamente kernels elementwise o reducciones, de huella fija y liviana. -
cudaDeviceSynchronizeal final de la ventana activa: el profiler fuerza un sincronizado para volcar eventos. Si esa sincronización cubre microsegundos de trabajo GPU pero milisegundos de sincronización, es otro signo de overhead-bound. -
cudaMemcpyAsyncque no escribiste: muchas veces es un Device-to-Device copy que prepara la memoria antes de un GEMM con epílogo (por ejemplo bias add). Inductor o cuBLAS pueden hacer ese memcpy implícito.
Variabilidad en tiempos de kernel
Incluso con el mismo kernel y los mismos datos, verás variabilidad entre invocaciones. Causas comunes:
- Boost de reloj en GPUs después de periodos de inactividad.
- Thermal throttling o gestión de energía.
- Tareas del driver o housekeeping del sistema.
La moraleja: no te fíes solo de la media. La traza te muestra la distribución real.
Qué cambia con torch.compile y qué no
torch.compile (por ejemplo con el backend predeterminado TorchInductor) captura fragmentos de ejecución y genera código: pipeline general resumido:
- TorchDynamo captura ejecución Python en un grafo FX.
- AOTAutograd prepara forward/backward cuando hay gradientes.
- Inductor baja el grafo a código GPU/CPU optimizado.
En la traza verás filas como "Torch-Compiled Region" y llamadas a TorchDynamo Cache Lookup en cada paso: Dynamo revisa si los shapes y metadatos casan con la cache. Ese chequeo se paga en cada llamada.
Un detalle importante observado en el ejemplo:
- Inductor puede fusionar al nivel del dispatcher: tu
torch.add(torch.matmul(x, w), b)puede convertirse enaten::addmm(b, x, w). - Sin embargo, eso no garantiza un único kernel GPU nuevo. En el caso 4096x4096 bf16, Inductor escogió usar cuBLAS y arregló la bias add como un epílogo de GEMM, acompañada de un
Memcpy DtoDpara preparar el buffer de salida.
Resultado práctico: compile puede aumentar el costo por llamada en la CPU (Dynamo + AOTAutograd + Inductor) para operaciones muy pequeñas. Para modelos con docenas o cientos de ops, ese overhead se amortiza. Para una operación aislada, es una tasa que te come tiempo.
Diagnóstico rápido: señales y cómo interpretarlas
Usa esto como chequeo cuando mires una traza:
-
Self CPU time total mucho mayor que Self CUDA time total (CPU ms vs GPU us): overhead-bound. Haz el trabajo más grande o fusiona.
-
Self CPU ≈ Self CUDA en ms: compute-bound. GPU es el cuello de botella.
-
Un evento domina
cuda_time_total: ahí empieza la optimización. -
Un evento con muchas llamadas (# of Calls grande): potencial para batching o fusión.
-
cudaOccupancyMaxActiveBlocksPerMultiprocessoraparece: kernel pesado adaptativo (GEMM, conv). -
cudaMemcpyAsyncinesperado antes de GEMM: copia DtoD para epílogo de addmm. -
Actividad Buffer Request en GPU lane: el profiler o la runtime está pidiendo/llenando buffers y puede explicar offsets entre lanes.
-
Variación en tiempos del mismo kernel entre steps: lee la traza completa, no solo la media.
Recomendaciones prácticas rápidas
-
Empieza perfilando con casos simples y viste cómo cambian las métricas al aumentar tamaño y batching.
-
Usa
warmup(tanto fuera como dentro del profiler) para eliminar costos de cold-start. -
Si el CPU domina, intenta fusionar operadores o mover trabajo al GPU en bloques más grandes.
-
Revisa la traza completa en Perfetto para encontrar gaps y memcpy implícitos; la tabla por sí sola no lo contará todo.
-
Cuando uses
torch.compile, inspecciona el árbol CPU: Dynamo cache hits, AOT wrapper y elCall CompiledFxGraphte dirán cuánto se está pagando en CPU.
Al terminar este primer ejercicio con un matmul + bias tienes ya varios lentes para leer trazas más complejas: distinguir overhead de cómputo, identificar planificación cuBLAS, entender copias implícitas y ver qué hace realmente torch.compile. En los siguientes posts subiremos la complejidad: primero nn.Linear y un MLP, y luego modelos grandes con transformers. La idea es construir la intuición paso a paso para que cuando abras una traza no sea terror sino curiosidad de detective.