Cómo domar Blackwell sin volverse loco con CUDA
Cuando se trata de escribir kernels GPU, la mayoría de nosotros inmediatamente想象amos páginas interminables de código CUDA C++, manipulando punteros y la agonizante depuración de memory leaks o condiciones de carrera. La situación se vuelve aún más compleja cuando monstruos como la arquitectura NVIDIA Blackwell (sm_100a) llegan al mercado. Hay tanta asincronía y nuevos motores de datos internos que los viejos enfoques de "simplemente lanza un hilo a un núcleo" ya no logran extraer ni la mitad del rendimiento.
Recientemente me topé con el proyecto modern-gpu-programming-for-mlsys del equipo MLC-AI. Básicamente, no es solo un repositorio, sino un tutorial interactivo vivo sobre programación GPU moderna. Los tipos fueron a fondo: te enseñan a escribir kernels SOTA (state-of-the-art) usando Python y un nuevo DSL llamado TIRx.
Por qué esto no es solo otro tutorial
Por lo general, los cursos de GPU se quedan atascados en la multiplicación de matrices clásica de la década de 2010. Aquí, el enfoque se desplaza hacia el hardware del mañana. El centro de atención es la arquitectura Blackwell, su jerarquía de memoria, núcleos tensor y mecanismos de movimiento de datos asíncronos.
Los autores proponen un camino desde la comprensión del hardware hasta escribir kernels que realmente funcionan en sistemas modernos de ML. En lugar de luchar con C++ de bajo nivel, escribes en TIRx (Tensor IR next). Esta es una extensión de Apache TVM que te permite describir la lógica del kernel en Python mientras mantienes el control sobre registros, caché y motores tensor.
Qué hay dentro de este tutorial
Todo el contenido se divide en bloques lógicos que elevan gradualmente la barra de dificultad.
Primero, cubren el modelo de ejecución y memoria. No hay definiciones aburridas aquí, pero sí mucha práctica trabajando con TMA (Tensor Memory Accelerator) y entendiendo cómo los datos realmente llegan a las unidades de cómputo. Curiosamente, los autores se detienen en detalle en el modelo Roofline—ayuda a entender en qué está bottlenecked tu código: el ancho de banda de memoria o la potencia bruta de los núcleos.
Luego viene la inmersión profunda en TIRx. Aprendes a trabajar con tiles (TileLayout), ejes y el mecanismo de swizzle (reordenamiento inteligente de datos en memoria para evitar conflictos de bancos). Al final, obtienes un kernel GEMM (multiplicación de matrices) funcional.
Después viene la "carne" real. Los autores muestran cómo convertir un GEMM regular en un monstruo de alto rendimiento. Las tuberías TMA, la especialización de warps y la agrupación de CTAs en clusters entran en juego. Si estos términos suenan como encantamientos, el tutorial los desglosará ordenadamente.
En el final, creas Flash Attention 4. Esto es exactamente en lo que funcionan los modelos grandes de lenguaje modernos. El kernel incluye dos etapas MMA (Matrix-Matrix Accumulation), softmax en línea y soporte para GQA (Grouped Query Attention).
Stack técnico y características
El proyecto está estrechamente acoplado con el ecosistema Apache TVM. Para ejecutar los ejemplos, necesitarás:
- Biblioteca
apache-tvm(TIRx viene incluido). - PyTorch para verificar resultados y preparar datos.
- Una GPU Blackwell (ej., B200) si quieres ejecutar el código en hardware real.
Si no tienes una Blackwell a mano (lo cual es lógico para la mayoría de nosotros ahora mismo), el repositorio sigue siendo extremadamente útil como base teórica. La lógica de optimización de memoria y programación de cómputos se aplica también a arquitecturas más antiguas, aunque algunas características como TMA son específicas de las generaciones NVIDIA más nuevas.
Para construir el libro localmente y leerlo en un formato conveniente, Sphinx estándar es suficiente:
pip install -r requirements-docs.txt
sphinx-build -b html . _build/html
¿Vale la pena tu tiempo?
A menudo me encuentro con desarrolladores de infraestructura de ML que tienen miedo de sumergirse en la escritura de kernels, considerándola prerrogativa de los "dioses del C++". Este proyecto demuestra que con las abstracciones correctas (como TIRx) y la comprensión de la arquitectura, puedes escribir código de clase mundial mientras permaneces en el ecosistema Python.
El proyecto será especialmente útil para quienes trabajan en optimizar redes neuronales "bajo el capó" o quieren entender por qué Flash Attention es más rápido que el enfoque habitual. Incluso si no planeas lanzar tu propia biblioteca Blackwell mañana, revisar cómo funcionan la memoria tensor y las colas asíncronas realmente aclara la mente y cambia tu perspectiva sobre la computación de alto rendimiento.
El único inconveniente es que la barrera de entrada sigue siendo alta. Esto no es "GPU para tontos", sino más bien "GPU para quienes quieren convertirse en profesionales". Pero si estás listo para descubrir cómo los bytes realmente se mueven dentro de una tarjeta gráfica, es difícil encontrar mejor material sobre el tema ahora mismo.
Proyectos relacionados