Cómo hacer que una GPU de AMD corra a velocidad de ‘brrr’

• HipKittens es una colección de primitivas de programación diseñada para extraer el rendimiento potencial de las GPU de AMD, optimizando el acceso a memoria, la planificación y la reutilización de caché
• La GPU AMD MI355X tiene una estructura de 256 unidades de cómputo y 8 chiplets (XCD), y ofrece un archivo de registros grande e instrucciones de núcleo matricial detalladas
• A diferencia de NVIDIA, AMD no tiene reasignación de registros, instrucciones matriciales asíncronas ni mbarrier, por lo que en lugar de wave specialization resultan más efectivos los esquemas de planificación de 8-wave ping-pong y 4-wave interleave
• HipKittens mejora la localidad de caché L2 y LLC mediante una planificación de grid consciente de chiplets, logrando mejoras en ancho de banda máximo y TFLOPS en operaciones GEMM y Attention
• Este enfoque compensa la falta de madurez del software en el ecosistema de GPU de AMD y sienta una base para ampliar la escalabilidad del cómputo de IA sobre hardware diverso

---

Arquitectura y características de rendimiento de las GPU AMD CDNA

• La GPU AMD MI355X incluye 256 unidades de cómputo (CU), y cada CU está compuesta por 4 SIMD
  • Un SIMD ejecuta una wave de 64 hilos, en contraste con el warp de 32 hilos de NVIDIA
• La MI355X tiene un 70% del SRAM de la B200 (165 KB) y carece de funciones como instrucciones asíncronas de multiplicación matricial, reasignación de registros, aceleración de memoria tensorial y mbarrier
• En cambio, ofrece un archivo de registros 2 veces más grande y 60% más procesadores (256 CU frente a 160 SM)
  • Soporta instrucciones de núcleo matricial pequeñas y detalladas, y existe una función de carga directa de memoria global a memoria compartida (similar a TMA)
• AMD adopta una arquitectura de chiplets compuesta por 8 chiplets (XCD); cada XCD tiene su propia caché L2 independiente, y por encima existe una caché LLC
• Según la tabla, la MI355X ofrece un rendimiento de BF16 2.5 PFLOPs, MXFP8 5.0 PFLOPs y MXFP6 10.1 PFLOPs, junto con 288 GB de capacidad de memoria y 8 TB/s de ancho de banda

Retos de diseño de kernels para AMD

• Optimización de acceso a memoria: debido a las limitaciones del compilador HIPCC y al comportamiento de I/O no público, es importante diseñar la disposición de datos y los patrones de swizzle
• Planificación dentro del procesador: AMD debe aprovechar el archivo de registros y las instrucciones matriciales pequeñas en lugar de depender de la memoria compartida
• Planificación entre procesadores: debido a la estructura basada en chiplets, es necesario distribuir el trabajo considerando efectos NUMA a nivel de caché

Patrones de acceso a memoria de HipKittens

• HipKittens (HK) usa tiles como unidad básica de datos y ofrece funciones de operación similares a PyTorch
  • Un tile se define por tipo de dato, tamaño y layout, y mediante metaprogramación con plantillas de C++ se adapta a distintas entradas
• Planificación de registros: como HIPCC no puede usar ciertos registros como entradas MFMA, HK ofrece una función explícita de fijación de registros
  • Esto permite que el desarrollador especifique directamente los registros para escribir kernels de máximo rendimiento
• Layout de registros: en AMD, el layout cambia según el tipo de dato y la forma de la matriz, así que no existe un único patrón de swizzle
  • Por ejemplo, un tile bf16 de 16×16 y un tile bf16 de 16×32 requieren patrones de swizzle distintos
• Estructura de fases de instrucciones: las instrucciones de memoria compartida de AMD tienen grupos de fases no contiguos y poca documentación interna
  • HK ofrece un solver obtenido por ingeniería inversa para manejar esto
• Generación de direcciones: AMD soporta cargas asíncronas de HBM a memoria compartida y realiza optimización con swizzle de direcciones HBM

Planificación dentro del procesador: patrones de wave

• Wave specialization es efectiva en NVIDIA, pero en AMD degrada el rendimiento por la ausencia de reasignación de registros
  • Las waves productoras ocupan registros innecesarios, y las consumidoras sufren spill por falta de registros
• En los experimentos de HK, wave specialization provoca en AMD menor intensidad aritmética y cuellos de botella de memoria
  • Ejemplo: en GEMM, la configuración HK 0/8 alcanza 1605 TFLOPs, mientras CUTLASS llega a 1570 TFLOPs
• Patrones de planificación alternativos
  • 8-wave ping-pong: dos waves ejecutan alternadamente clústeres de memoria/cómputo
  • 4-wave interleave: una wave intercala finamente memoria y cómputo
  • El esquema de 8-wave produce código más conciso, mientras que 4-wave es más detallado pero más extenso
  • En GEMM y Attention Forward, 8-wave logra rendimiento de nivel SoTA

Planificación entre procesadores: enfoque consciente de chiplets

• La AMD MI355X tiene 8 chiplets XCD, y cada chiplet cuenta con una caché L2 independiente
  • Como los bloques de hilos se asignan a los chiplets en esquema round-robin, el orden del grid afecta directamente la eficiencia de reutilización de caché
• Una disposición simple en row-major tiene una tasa baja de reutilización de caché L2, lo que provoca pérdida de ancho de banda
  • Ejemplo: L2 55%, LLC 95%, 15.1 TB/s, 1113 TFLOPs
• HK introduce una planificación de grid consciente de chiplets para aprovechar al mismo tiempo la localidad de caché L2 y LLC
  • Agrupa bloques de hilos por regiones adyacentes de la matriz de salida para maximizar la reutilización de datos de entrada

Ejemplos de kernels reales

• El hot loop de los kernels de Attention Forward y BF16 GEMM usa el esquema de 8-wave ping-pong de HK
  • Cada loop ejecuta alternadamente clústeres de Compute–Memory y se sincroniza con una schedule barrier
  • En el ejemplo de código, operaciones de HK como mma_AtB, load, exp2, col_sum se usan repetidamente

Conclusión: AMD en la era de la IA multi-silicon

• HipKittens logra rendimiento competitivo en AMD CDNA3 y CDNA4
  • Hay tres claves: acceso a memoria optimizado, planificación de waves centrada en AMD y planificación de grid consciente de chiplets
• Los kernels de HK alcanzan el mejor rendimiento dentro del ecosistema AMD y compiten incluso con kernels de NVIDIA Blackwell
• Para la diversidad del cómputo de IA, hace falta ampliar la accesibilidad a las GPU de AMD, e HipKittens aporta la base de software clave para ello
• La mejora de la planificación de registros en HIPCC se señala como un área importante de avance futuro

[GN⁺]

Opiniones en Hacker News

• Recomiendo revisar la discusión relacionada sobre HipKittens
• También está el artículo HipKittens: Fast and furious AMD kernels, que trata la misma investigación. Tiene comentarios de George Hotz y empleados de AMD
• Me alegra que la academia esté abordando este tipo de problemas, pero al final creo que este es un problema que AMD debe resolver internamente
  • Yo creo que lo mejor es que las empresas de hardware se dediquen solo al hardware. Así los incentivos se mantienen puros. Incluso si el rendimiento cae 20%, me parece preferible
  • Totalmente de acuerdo. AMD fue posponiendo este problema durante 10 años y apenas ahora intenta ponerse al día. El hardware es excelente, pero la falta de capacidad para desarrollar firmware le impide aprovechar todo su potencial
  • Pero este equipo de investigación también ha creado software similar para GPU de Nvidia. Parece que son investigadores muy capaces aplicando su especialidad
  • Hasta donde sé, AMD ya está abordando este problema en varios niveles y también está colaborando con tinycorp
• El artículo deja la impresión de que la optimización es difícil por la complejidad arquitectónica de las GPU de AMD. Pero a largo plazo, el enfoque de AMD podría escalar mejor. Mientras Nvidia usa 2 chiplets, AMD tiene una estructura de 8 chiplets, así que hay problemas de localidad de memoria. Como en el futuro probablemente habrá aún más chiplets, la experiencia de lidiar con esta complejidad hoy podría ayudar a largo plazo
  • AMD no necesita warp specialization para lograr alto rendimiento, así que programar es más simple
• Muchos desarrolladores han intentado hacer que las GPU de AMD “go brrr” para desarrolladores comunes, pero han fracasado. No entiendo por qué AMD no resuelve por sí misma los problemas de software. Ya tiene suficiente dinero, así que no contratar desarrolladores ya no es excusa. Sus GPU para datacenter no están mal, pero para que una persona haga experimentos de ML o IA, Nvidia sigue siendo muchísimo mejor. Siento que mi RTX 3090 de hace 5 años sigue siendo mejor que cualquier GPU de consumo de AMD lanzada hasta ahora
  • La experiencia para desarrolladores en AMD es terrible. Ni siquiera aceptan reportes de bugs por crashes de drivers
  • Hace poco cambié mi servidor de inferencia de una NVidia 5090 a dos AMD R9700 32GB, y la experiencia fue totalmente positiva. Funcionó de inmediato en el kernel de Fedora sin configurar DKMS, y conectar contenedores con ROCm fue fácil. Solo tuve que cambiar la configuración de Ollama y Storyteller. Fue una experiencia mucho más agradable que con CUDA
  • Nvidia incluso mantiene directamente un fork de Unreal Engine. AMD ni siquiera compite en ese nivel
  • Nvidia es la única empresa de hardware que ofrece a los ingenieros de software una compensación competitiva. En AMD todavía persiste una cultura que no ve el software como “trabajo de verdad”, y es difícil cambiar esa inercia
• Mojo tenía la idea de mejorar la experiencia de desarrollo (devX) en GPU de AMD; me da curiosidad saber cómo va eso
• No entiendo por qué AMD no invierte miles de millones de dólares en mejorar el software. Nvidia es la empresa más valiosa del mundo, y AMD es su único competidor real
  • AMD sí lo está intentando, pero creo que es difícil transformar una cultura organizacional de renovar hardware cada año en una cultura centrada en software. El software no genera ingresos inmediatos como el hardware, así que la gerencia tiende a darle menos prioridad. Además, que proveedores externos aporten código open source puede verse bien a corto plazo, pero perjudica la calidad a largo plazo. Y si te pierdes una sola tendencia de hardware, corres el riesgo de quedar atrás frente a la competencia
  • He trabajado con varios fabricantes de GPU, y solo Nvidia ve el software como un activo (asset) e invierte en él. Las demás compañías lo ven solo como un costo
• En lo personal no me gusta mucho el meme de “go brr”, pero me da risa verlo usado en lugares como Stanford
  • De hecho ya lo habían usado hace un año, durante la presentación de ThunderKittens
  • Si ese tipo de memes ya aparece en canales oficiales de una universidad, quizá sea una señal de que la moda ya pasó
• El proyecto en sí es excelente, pero sigo preguntándome por qué AMD no hace esto directamente. Parece que AMD todavía no entiende la importancia de una pila de software madura. Necesita una pila unificada como CUDA, que funcione en todas las tarjetas. Antes creía que AMD eventualmente se pondría al día, pero ahora casi ya perdí la esperanza
• El proyecto está bien, pero el artículo en sí se siente escrito de forma extraña
  • Está demasiado raro. Parece que dependieron demasiado de IA, o que intentaron imitar un estilo de escritura de IA. Se repiten frases como “revisa la parte uno” o “cómo hacer que las GPU de AMD hagan go brr”. También decepciona que explicaran partes técnicas que deberían ir en gráficos con 100 líneas de código