¿Qué es exactamente un kernel?

Dos números en una GPU

Ya sabes cómo hacer c = a + b en Python. Lo has hecho mil veces. Aquí vamos a hablar de ello con tensores de PyTorch. Un tensor es simplemente un array de números. Ponerlo en una GPU significa que ese array vive en la memoria de la GPU en lugar de la memoria normal de la CPU. Cuando a y b son dos tensores en una GPU, esa línea termina lo suficientemente rápido como para que nunca pienses en ello.

Ahora redúcelo. Supón que a y b son dos flotantes individuales, ambos en la GPU. La misma línea: c = a + b. ¿Qué se ejecuta realmente en el chip?

La respuesta es un kernel. En este contexto, un kernel es un pequeño programa que la GPU ejecuta sobre algún fragmento de datos. No es el kernel del sistema operativo con el que arranca tu portátil. No es el kernel matemático de un libro de álgebra lineal. La palabra se reutiliza mucho, y no es tu culpa. En el mundo de las GPU, un kernel significa simplemente: una pequeña función a la que la GPU recibe la orden de ejecutar, en ese momento, en paralelo, sobre los datos que le entregaste.

Al final de esto, serás capaz de mirar un fragmento de código de PyTorch y contar cuántos kernels ejecutará la GPU. Suena a truco pequeño, y lo es, pero también es el primer paso para sacudirse esa sensación de "GPU como caja negra", esa en la que tu modelo va lento y no sabes por qué. Toda pregunta que puedas hacer sobre el rendimiento de la GPU termina volviendo a "qué kernels se ejecutaron y qué estaban haciendo". Así que aquí empezamos.

Tu primer kernel

Vamos a hacer que a y b sean un poco más grandes. Tensores de longitud 8 esta vez. Sigue siendo una línea de PyTorch: c = a + b.

Cuando ejecutas esto, tu CPU (la máquina que ejecuta el Python real) le dice a la GPU: oye, ejecuta este programa sobre estos datos. Esa instrucción se llama un lanzamiento. Lo que se lanza es un kernel: un programa, listo para ejecutarse. Los lanzamientos en sí son baratos, microsegundos cada uno. Lo que está alrededor del lanzamiento (los datos que van a la GPU, los resultados que vuelven) es donde reside el coste real, y eso es lo que iremos contando.

Dentro del kernel, el trabajo real lo realizan pequeños trabajadores llamados hilos. Una GPU tiene miles de ellos disponibles. Para nuestra suma de longitud 8, 8 hilos recogen el trabajo: el hilo 0 maneja el elemento 0, el hilo 1 maneja el elemento 1, y así hasta el hilo 7. Cada hilo ejecuta el mismo pequeño programa: leer un elemento de a, leer el elemento correspondiente de b, sumarlos, escribir el resultado en c.

(En la práctica, la GPU lanza hilos en grupos de tamaño fijo llamados warps, siempre 32 hilos en las tarjetas NVIDIA, y enmascara los sobrantes cuando tu array no se divide uniformemente. Puedes ignorarlo por ahora.)

Así que tenemos una línea de PyTorch, un lanzamiento, un kernel, 8 hilos haciendo 8 sumas. Ahora contemos qué cruzó realmente el chip. Para hacer la suma, cada hilo necesita su elemento de a y su elemento de b. Eso son 8 lecturas de a y 8 lecturas de b. Luego, cada hilo escribe su resultado en c. Eso son 8 escrituras.

Esas lecturas y escrituras van a la gran memoria que está justo al lado del chip de la GPU. En tarjetas de centro de datos (A100, H100) esa memoria se llama HBM (high-bandwidth memory). En tarjetas de consumo (RTX 4090) y las T4 estilo Colab en las que la gente probablemente probará esto, se llama GDDR. De cualquier manera, es memoria rápida junto al chip, y la llamaremos simplemente memoria de la GPU. Es rápida, pero no es gratis, y cada viaje a ella tiene un coste.

Un kernel = un lanzamiento = un pase sobre los datos. Sea lo que sea que el kernel haga en su interior, las lecturas y escrituras en sus bordes (los viajes a la memoria de la GPU para obtener entradas, el viaje de vuelta para escribir salidas) son la parte que cuesta. Esa es toda la forma.

Nada de esto cambia cuando los tensores se hacen más grandes. La misma línea de PyTorch, el mismo kernel, solo que más hilos. Si a y b tienen un millón de elementos cada uno, la GPU lanza el mismo kernel con un enjambre más grande de hilos. Las matemáticas escalan, los bytes movidos escalan, el modelo mental no. Una línea, un kernel.

Lo que ocurre entre dos operaciones

1c = (a + b).relu()

Conoces Python lo suficiente como para saber que esto son dos operaciones, una suma y luego una relu, encadenadas. En un intérprete, eso son dos llamadas a función. En una GPU, en PyTorch eager, eso son dos lanzamientos de kernel: uno para la suma, otro para la relu. Hasta ahora, nada sorprendente.

Lo realmente interesante es lo que ocurre entre los dos kernels.

Cuando la suma termina, tiene que poner su resultado en algún sitio. Ese sitio es la memoria de la GPU. La suma escribe un array intermedio completo (llámalo tmp) en la memoria. Luego, un momento después, se lanza la relu, y su primer trabajo es leer ese mismo array tmp de vuelta de la memoria. Lee todo el array, aplica relu a cada elemento, escribe el resultado en c.

Cuenta el tráfico de memoria para esos dos kernels:

• La suma: lee a, lee b, escribe tmp. Tres transferencias del tamaño del array.
• La relu: lee tmp, escribe c. Dos más.

Cinco transferencias del tamaño del array en total. Compara eso con la suma de longitud 8 por sí sola de la sección anterior, que eran tres. Añadir .relu() a la cadena no solo te costó el cómputo de la relu. Te costó un viaje completo de ida y vuelta del array a través de la memoria de la GPU, porque tmp tuvo que ser escrito solo para que el siguiente kernel pudiera leerlo de vuelta.

Nada se almacenó en caché. tmp no pudo quedarse en un registro o en una caché local rápida. Salió a la memoria de la GPU (la lenta y lejana) y volvió inmediatamente. Los dos kernels son desconocidos el uno para el otro. Tienen que pasar los datos a través del único medio que ambos saben usar: la memoria de la GPU.

¿Por qué PyTorch lo hace así? Porque en modo eager, cuando escribes a + b, PyTorch lo ejecuta ahora. No sabe que ibas a llamar a .relu() a continuación. Cada operación se envía a ejecutar en el momento en que su línea de Python se ejecuta. No hay plan, no hay previsión. Cada operación se sostiene por sí misma, produce un array real, y lo pasa a lo que sigue a través de la memoria.

Este es el patrón que debes retener. Cada valor intermedio en tu código de PyTorch se escribe físicamente en la memoria de la GPU y es leído de vuelta por la siguiente operación. Cada uno. Eso es lo que realmente mide el "conteo de kernels". Cada kernel extra es otro viaje de ida y vuelta que tus datos tienen que hacer a través de la memoria de la GPU.

Fusión: dos operaciones, un kernel

Imagina un kernel que hace todo esto de una sola vez: lee su elemento de a, lee su elemento de b, los suma, aplica relu al resultado (todo dentro del kernel, en un pequeño espacio de trabajo privado por hilo que nunca sale del chip), y solo entonces escribe el valor final en c. El valor intermedio (a + b) sigue existiendo, pero solo dentro del kernel, en el espacio privado de cada hilo. Nunca se escribe en la memoria de la GPU. tmp, como array real, no existe en absoluto.

Cuenta las transferencias ahora. Lecturas de a: 1 por elemento. Lecturas de b: 1 por elemento. Escrituras de c: 1 por elemento. Tres transferencias del tamaño del array. Las mismas matemáticas que la versión de dos kernels, pero dos viajes de ida y vuelta menos.

Con longitud 8, esto es un error de redondeo. A nadie le importa. Con longitud 1 millón, o 100 millones, esos viajes de ida y vuelta extra se convierten en una gran parte del tiempo de ejecución, y el reloj de pared lo refleja. Por qué el tráfico de memoria termina dominando así es el tema completo del Artículo 2 de esta serie, así que dejaré el "por qué" aquí. El punto ahora es solo: mismas matemáticas, menos viajes, más rápido en la práctica.

Ese truco, combinar operaciones que habrían sido kernels separados en un solo kernel para que el valor intermedio nunca tenga que visitar la memoria de la GPU, tiene un nombre. Se llama fusión. Esa es toda la palabra. Esa es toda la idea.

Ahora la parte incómoda. Escribir ese kernel combinado a mano parece fácil para suma + relu. Dos operaciones. Una línea de "cómputo" en medio. Pero el código real de PyTorch tiene docenas de operaciones encadenadas, cada una con sus propias formas, tipos de dato y reglas de broadcasting. Escribir un kernel fusionado que maneje todo eso correctamente es un trabajo de ingeniería real. Generalmente no escribirías a mano estos kernels elementales de rutina.

Buenas noticias: PyTorch incluye una herramienta que hace la reescritura por ti, para exactamente este tipo de casos, automáticamente. Se llama torch.compile.

Casi con seguridad has visto torch.compile usado así:

1model = torch.compile(model)

Una línea. Alguien en internet te dijo que hace las cosas más rápidas. Esto es lo que realmente hace, en español sencillo: en lugar de ejecutar tus operaciones una por una como hace el modo eager, torch.compile captura las operaciones con tensores que realiza tu función, busca oportunidades para combinarlas y genera código optimizado. Las llamadas posteriores que coincidan con las mismas suposiciones pueden reutilizar ese código.

La fusión que hicimos en papel arriba (suma y relu, compartiendo un kernel, tmp sin tocar la memoria) es exactamente el tipo de cosa que torch.compile hará automáticamente en tu código, siempre que las operaciones sean lo suficientemente simples. Cuando la gente dice que torch.compile "hace PyTorch más rápido", una gran parte de eso es lo que quieren decir.

Para los casos que torch.compile no puede fusionar por sí solo (operaciones personalizadas que no reconoce, reducciones inusuales, layouts de memoria extraños), alguien tiene que escribir un kernel a mano. Para eso sirven herramientas como Triton y CUDA. Un artículo aparte.

Compruébalo tú mismo

Todo lo anterior ha sido contar kernels en papel. Es hora de contarlos en una GPU real. Si tienes alguna máquina con una GPU CUDA a mano (una estación de trabajo, un notebook de Colab, una instancia en la nube), puedes ejecutar esto tú mismo en unos minutos.

La herramienta es torch.profiler. Viene integrada en PyTorch. Todo lo que hace es registrar lo que la GPU realmente hizo mientras tu código se ejecutaba, y devolverte una tabla que puedes leer.

Paso 1: la versión eager.

Envuelve la línea de dos operaciones en una función para tener algo que llamar:

1import torch
2from torch.profiler import profile, ProfilerActivity
3
4def add_relu(a, b):
5    return (a + b).relu()
6
7a = torch.randn(1_000_000, device="cuda")
8b = torch.randn(1_000_000, device="cuda")
9
10with profile(activities=[ProfilerActivity.CUDA]) as prof:
11    add_relu(a, b)
12    torch.cuda.synchronize()
13
14print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

torch.cuda.synchronize() está ahí solo para asegurarse de que la GPU ha terminado antes de que leamos los tiempos. El trabajo de la GPU se ejecuta de forma asíncrona, y sin la sincronización a veces medirías la sobrecarga del lanzamiento en lugar del trabajo real del kernel.

Paso 2: lee la salida.

Tu salida real del profiler tendrá más filas de las que esperas. Se mezclarán un montón de líneas de asignación de memoria y contabilidad interna de PyTorch. Las filas que nos interesan son los kernels CUDA, las funciones reales que ejecutó la GPU. Busca filas con kernel en el nombre. Las dos se verán más o menos así:

1vectorized_elementwise_kernel<...BinaryFunctor_add...>    1    12us
2vectorized_elementwise_kernel<...threshold_kernel...>     1     9us

Los nombres exactos de las plantillas cambian entre versiones de PyTorch (relu suele aparecer como threshold porque esa es la operación subyacente, y add a veces como CUDAFunctor_add). No intentes analizar todo. Solo cuenta las filas. Dos filas. Dos kernels. Uno para la suma, otro para la relu. Exactamente lo que dijimos que ocurriría hace una sección.

Paso 3: la versión compilada.

Un cambio de línea. Envuelve la función con torch.compile:

1compiled = torch.compile(add_relu)

Ahora, aquí hay un detalle que vale la pena saber antes de ejecutarlo. La primera llamada a una función compilada con torch.compile es lenta, a veces extrañamente lenta, porque es entonces cuando torch.compile hace su trabajo: analizar tu código, decidir qué fusionar, generar el kernel fusionado. Si perfilas la primera llamada, estás midiendo el paso de compilación, no el kernel. Así que el patrón es: llámala una vez para calentarla, descarta el resultado, y luego perfilala.

1compiled(a, b)                # calentamiento, descartar
2torch.cuda.synchronize()
3
4with profile(activities=[ProfilerActivity.CUDA]) as prof:
5    compiled(a, b)
6    torch.cuda.synchronize()
7
8print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

Paso 4: lee la salida otra vez.

Esta vez:

1triton_poi_fused_add_relu_0    1    14us

Una fila. Un kernel. El nombre incluso te dice lo que hizo: suma y relu fusionadas. Las mismas matemáticas que antes, un lanzamiento en lugar de dos.

Acabas de hacer lo que este artículo ha estado explicando, en una frase: le pediste a PyTorch que combinara dos operaciones en un kernel, miraste el profiler y confirmaste que el conteo pasó de dos a uno. Fusión, en el mundo real, en tu máquina.

Si quieres verlo de forma más dramática, pruébalo con diferentes tamaños de tensor. Con longitud 100, ambas versiones se ejecutan tan rápido que la diferencia se pierde en el ruido. Con longitud 10 millones o 100 millones, la versión compilada empieza a destacar claramente, porque el viaje de ida y vuelta que eliminamos es una parte real del trabajo a esa escala.

Contar kernels ya no es un consejo abstracto. Tienes una forma de comprobarlo.

Resumamos

Aquí está todo de un tirón.

Tu código de PyTorch, cuando se ejecuta en una GPU, se convierte en una secuencia de kernels. Cada kernel es un lanzamiento, un pase sobre tus datos, un viaje de ida y vuelta a través de la memoria de la GPU para obtener entradas y escribir salidas. Las operaciones simples se convierten en un kernel. Las cadenas de operaciones se convierten en un kernel por operación por defecto, con valores intermedios haciendo viajes de ida y vuelta a través de la memoria entre ellos. torch.compile puede fusionar cadenas simples por ti para que esos valores intermedios nunca toquen la memoria. Menos kernels suele significar menos tráfico de memoria. Y menos tráfico de memoria suele significar más rapidez.