La mejora explosiva del rendimiento de las GPU

• En un contexto donde el costo del cómputo de IA sigue creciendo, Hazy Research resume que la clave para optimizar el rendimiento de la GPU es mantener siempre ocupados los tensor cores de la NVIDIA H100
• La H100 alcanza 989 TFLOPs en multiplicación de matrices de media precisión, pero en operaciones generales se queda en unos 60 TFLOPs, por lo que en cuanto los tensor cores se detienen, la utilización cae con fuerza
• Para acercarse al rendimiento máximo hay que manejar en conjunto WGMMA, la disposición de la shared memory, la generación de direcciones y la occupancy; sin wgmma.mma_async, en microbenchmarks se queda en alrededor del 63% del pico
• El DSL embebido en CUDA y publicado abiertamente, ThunderKittens, encapsula complejidades como el swizzling y el layout de registros mediante abstracciones de tile y vector, simplificando la escritura de kernels de la familia FlashAttention
• El kernel forward de FlashAttention-2 para H100 se escribe en unas 100 líneas y es cerca de un 30% más rápido que FlashAttention-2, mientras que el kernel de Based linear attention funciona a 215 TFLOPs

---

Condiciones que determinan el rendimiento de la H100

• La IA usa mucho cómputo, y Hazy Research ha trabajado en los últimos años en hacer que la IA use menos cómputo o que corra de forma más eficiente con el cómputo disponible
  • Como ejemplos de reducción de cómputo están Based, Monarch Mixer, H3, Hyena, S4
  • Como ejemplos de ejecución eficiente están FlashAttention, FlashAttention-2, FlashFFTConv
• El objetivo práctico es ordenar lo aprendido al acelerar GPUs y publicar el DSL embebido en CUDA ThunderKittens, que ayuda a escribir kernels rápidos
• En un sentido más amplio, también aborda cómo entender el hardware cambió la manera de ver el cómputo de IA

Arquitectura y cuellos de botella de la NVIDIA H100

• La discusión toma como base una GPU H100 SXM con la siguiente configuración
  • 80GB de HBM3, con ancho de banda de 3TB/s
  • 50MB de caché L2, con ancho de banda de 12TB/s, dividida en 2 secciones de 25MB a lo largo de la GPU y conectadas por un crossbar

• 132 SM

  • Cada SM tiene una caché L1 de 256KB que incluye hasta 227KB de shared memory, con un ancho de banda combinado de unos 33TB/s
  • El nuevo hardware de Hopper, Tensor Memory Accelerator (TMA), se encarga de la generación asíncrona de direcciones y del fetch de memoria
  • Cada SM se compone de 4 quadrants, y cada quadrant tiene warp scheduler, 512 vector registers, tensor core para multiplicación de matrices e instrucciones especializadas en paralelo
  • Todo el cómputo ocurre en el SM, y la mayor parte se procesa en registers
  • La clave para sacar rendimiento en la H100 es mantener los tensor cores siempre alimentados
  • La H100 ofrece 989 TFLOPs en multiplicación de matrices de media precisión y unos 60 TFLOPs en las operaciones “restantes”
  • En los ciclos donde se usan tensor cores se alcanza al menos un 94% de utilización del hardware
  • En los ciclos donde no se usan tensor cores, la utilización se queda como máximo en 6%

WGMMA: una instrucción necesaria pero difícil

• La H100 incluye la instrucción warp group matrix multiply accumulate wgmma.mma_async
  • En PTX es wgmma.mma_async
  • En SASS aparece como HGMMA / IGMMA / QGMMA / BGMMA
• En GPUs anteriores, wmma.mma.sync y mma.sync funcionaban de forma síncrona: un warp de 32 threads cargaba datos en el tensor core y esperaba el resultado
• wgmma.mma_async hace que 128 threads contiguos se sincronicen de forma cooperativa a través de todos los quadrants del SM e inicien directamente desde shared memory una multiplicación de matrices asíncrona
  • Mientras la multiplicación avanza, los warps pueden hacer otras tareas en registers
  • El resultado puede esperarse en el momento deseado
• En microbenchmarks, estas instrucciones fueron necesarias para extraer toda la capacidad de cómputo de la H100
  • Si no se usan, se observó que la GPU se queda alrededor del 63% de utilización pico
  • Podría deberse a que el tensor core exige pipelines de hardware profundos incluso desde recursos locales
• La mayor dificultad está en la complejidad del memory layout
  • Un layout de shared memory sin swizzling tiene un coalescing muy pobre y exige mucho ancho de banda de L2
  • La documentación del layout con swizzling estaba mal, así que tomó tiempo entenderlo
  • El layout con swizzling parece funcionar solo con ciertas formas de matriz y no encaja bien con otras funciones de wgmma.mma_async
  • El hardware puede hacer transpose de submatrices en el camino al tensor core, pero solo cuando el layout no tiene swizzling
• En kernels como FlashAttention, TMA y la caché L2 son lo bastante rápidas como para ocultar parcialmente este problema
• Para usar por completo el hardware hay que coalescer requests de memoria y evitar bank conflicts, por lo que controlar el layout es importante

Shared memory y bank conflicts

• La latencia de acceso único a shared memory parece ser de unos 30 ciclos, tiempo durante el cual el tensor core del SM casi puede ejecutar dos multiplicaciones de matrices cuadradas de 32x32
• En trabajos anteriores como FlashAttention, el foco estaba principalmente en el cuello de botella HBM-SRAM, y en el pasado ese cuello sí era realmente importante
• A medida que la HBM se volvió más rápida y los tensor cores crecieron más deprisa que otras partes del chip, incluso la pequeña latencia de shared memory pasó a ser algo que hay que eliminar u ocultar
• La shared memory está dividida en 32 banks, así que si no se tiene cuidado aparecen bank conflicts
  • Si se solicitan al mismo tiempo varios fragmentos distintos del mismo memory bank, la solicitud se serializa
  • En la práctica, esto puede volver un kernel desproporcionadamente lento
  • El layout de registers que exigen las instrucciones WGMMA y MMA puede sufrir bank conflicts si se usa de forma ingenua
• La solución es reordenar la shared memory con distintos patrones de swizzling para evitar conflictos
• Siempre que sea posible, conviene evitar mover datos entre registers y shared memory; y cuando haga falta, es mejor usar hardware dedicado como WGMMA y TMA para hacer transferencias asíncronas
• Usar warps reales para mover datos de forma síncrona es lo más común, pero se acerca al peor fallback posible

Generación de direcciones y TMA

• En la H100, tanto los tensor cores como la memoria son tan rápidos que generar las direcciones de memoria a fetchar ya consume una parte considerable de los recursos del chip
  • Esto se nota aún más cuando se agregan patrones complejos de interleaving o swizzling
• El Tensor Memory Accelerator (TMA) de NVIDIA permite definir un layout tensorial multidimensional en global/shared memory, hacer fetch asíncrono de un subtile de ese tensor y disparar una barrier al completarse
• TMA reduce el costo de generar direcciones y también facilita construir pipelines
• Al igual que wgmma.mma_async, TMA se considera esencial para sacar el potencial de la H100
  • Según la experiencia del equipo, podría ser incluso más importante que WGMMA
  • Ahorra recursos de registers y de instruction dispatch
  • También puede hacer reducción asíncrona sobre global memory, lo que lo vuelve útil en kernels backward complejos
• Entender los modos de swizzling de TMA también requirió algo de reverse engineering, aunque fue menos doloroso que con WGMMA

El costo que oculta la occupancy

• En CUDA, occupancy se refiere a la cantidad de threads co-programados sobre el mismo hardware de ejecución
• El warp scheduler de cada quadrant del SM intenta emitir en cada ciclo una instrucción al warp que esté listo para recibirla
• La H100 depende menos de la occupancy que generaciones anteriores
  • Gracias a sus funciones asíncronas, incluso un solo instruction stream puede mantener ocupados al mismo tiempo el fetch de memoria, la multiplicación de matrices, la reducción en shared memory y las operaciones matemáticas en registers
• Aun así, la occupancy es muy útil para ocultar errores y costos de sincronización
  • Un pipeline perfectamente diseñado puede ser rápido incluso sin occupancy adicional
  • En la práctica, se observó que las GPUs de NVIDIA parecen estar diseñadas pensando en la occupancy
  • Como hay mucha sincronización y margen para errores, aumentar la occupancy suele mejorar la utilización real del hardware
• En la H100, la occupancy es útil en un nivel moderado, pero en la A100 y la RTX 4090 parece ser aún más importante
  • Se menciona como posible razón una mayor dependencia del instruction dispatch síncrono frente a la H100

ThunderKittens: un pequeño DSL dentro de CUDA

• ThunderKittens es un DSL embebido en CUDA creado para facilitar la escritura de kernels rápidos en la H100
• Al principio se hizo para uso interno del laboratorio y luego se publicó
• El nombre se eligió porque kittens suena tierno y les parecía divertido escribir kittens:: en el código
• ThunderKittens apunta a la simplicidad y ofrece cuatro tipos templated
  • Register tiles: tensor 2D sobre el register file
  • Register vectors: tensor 1D sobre el register file
  • Shared tiles: tensor 2D dentro de shared memory
  • Shared vectors: tensor 1D dentro de shared memory
• Los tiles se parametrizan por height, width y layout
• Los register vectors se parametrizan por length y layout, mientras que los shared vectors usan solo length
  • Los shared vectors por lo general no sufren bank conflicts
• Las operaciones disponibles manipulan tiles y vectors a nivel de warp o de warp group cooperativo
  • initializer: por ejemplo, poner en cero un shared vector
  • unary op: como exp
  • binary op: como mul
  • row/column op: como row_sum
• Como ThunderKittens está embebido en CUDA, a diferencia de bibliotecas como Triton, se explica que la abstracción falla de forma “graceful”
  • Si falta alguna función, se puede extender del modo deseado

Ejemplo de FlashAttention y rendimiento

• Como ejemplo de ThunderKittens se presenta un kernel forward simple de FlashAttention para RTX 4090
  • Solo cubre headdim=64
  • n debe ser múltiplo de 256
  • Está escrito en unas 60 líneas de código CUDA
  • La utilización del hardware es del 75%
  • La mayor parte de la complejidad está en el algoritmo mismo, no en el patrón de swizzling ni en el layout de registers
• También se escribió con ThunderKittens el forward pass de FlashAttention-2 para H100
  • ThunderKittens encapsula la complejidad de TMA, WGMMA, los modos de swizzling y los descriptores
  • El kernel tiene unas 100 líneas
  • En la H100 es cerca de 30% más rápido que FlashAttention-2
• ThunderKittens se describe como una especie de “mini-pytorch” para GPU que encapsula layouts e instrucciones y ofrece primitivas
• También se publican kernels para Based linear attention y para otras arquitecturas que se anunciarán después
  • El kernel de Based linear attention funciona a 215 TFLOPs
  • Si se considera el recompute propio del algoritmo, supera los 300 TFLOPs
  • En teoría, linear attention es más eficiente, pero en hardware real históricamente ha sido mucho menos eficiente
  • Se considera que este resultado podría ampliar el rango de aplicaciones de alto throughput

Pensar en torno al tile

• Se considera que ThunderKittens funcionó bien porque no intenta hacerlo todo
  • CUDA es mucho más expresivo que ThunderKittens
  • ThunderKittens es un DSL pequeño y simple
• La abstracción central es el small tile, y se plantea que esto encaja con la dirección hacia la que avanzan la IA y el hardware
• ThunderKittens no soporta dimensiones menores a 16
  • Se considera que el hardware tampoco favorece especialmente dimensiones tan pequeñas
  • Se plantea la pregunta de si una multiplicación de matrices menor que 16x16 realmente puede considerarse IA
• La visión de la era CPU, en la que una palabra de 32 bits se considera un register, no encaja con el hardware de IA
  • El vector register de 1024 bits en CUDA se ve como un paso en la dirección correcta
  • Aquí, el register contiene los datos de un tile de 16x16
• Como la IA sigue girando alrededor de matrix multiply, reduction y reshape, la abstracción de tile se considera adecuada tanto para la IA como para el hardware
• Hacia adelante, las ideas de IA deberán reorganizarse para mapear mejor al hardware
  • El tamaño del estado recurrente debe ser lo bastante grande como para caber en un SM
  • La densidad de cómputo no debe estar por debajo de lo que exige el hardware
  • Ajustar el diseño de IA a lo aprendido del hardware se plantea como una dirección importante a futuro

Plan de soporte para AMD

• El soporte de ThunderKittens para hardware de AMD llegará pronto

[GN⁺]

Opiniones de Hacker News

• Me parece interesante la pregunta: “si la multiplicación de matrices es menor que 16x16, ¿estás seguro de que eso realmente es IA?”
  Los requisitos del hardware para IA se están volviendo cada vez más claros. Las GPU se diseñaron originalmente para un propósito completamente distinto, pero se usaron para IA porque su hardware de multiplicación de matrices era bueno; y una “GPU de IA” podría prescindir de algunas funciones que sí están en una GPU real.
  También hay una tendencia hacia representaciones numéricas más cortas, como punto flotante de 16 bits, 8 bits, 2 bits y 1 bit, y algún día se definirá cuál es el punto adecuado. Este artículo muestra que el hardware que favorece tiles de 16x16 tiene bastante sentido. Es muy probable que alguien ya esté escribiendo algo así en VHDL, o que pronto lo haga.
  Al final, parece probable que aparezcan dispositivos más simples, menos generalistas y más baratos, que ejecuten solo operaciones de “IA” con la menor carga posible de hardware innecesario.

  • Las GPU ya han evolucionado hasta convertirse en máquinas de IA lo más depuradas posible. Al menos desde 2014, cuando se fundó Nervana, se decía que las GPU eran tecnología vieja y no eran adecuadas para IA, pero parece que no se esperaba que evolucionaran tan rápido hacia máquinas de IA.
  • Apple ya va en esta dirección desde hace varios años. La NPU en el die es completamente distinta de una GPU o una CPU[1]
    Nvidia probablemente también esté trabajando en esto, pero desde el punto de vista del negocio quizá le convenga más mantener un dispositivo que agrupe gaming/entretenimiento/criptomonedas/IA, es decir, en formato de tarjeta de video.
    [1] https://github.com/hollance/neural-engine/blob/master/docs/a...
  • La parte de “las mentiras de NVIDIA” muestra la profundidad de la competencia. Es difícil que un error en la documentación sea completamente casual, y como los diagramas son fáciles de robar o copiar, Nvidia podría haber visto cierta utilidad en dejarlos así a propósito.
    Me recuerda a la época en que la Nervana de Naveen Rao hacía un driver Nvidia Maxwell más rápido que el propio driver de Nvidia. No todos los errores de documentación de un producto que crece rápido son medidas contra la competencia, pero si consideramos cuánto tardaron los investigadores en hacer ingeniería inversa de wgmma y hasta la situación política entre EE. UU. y China alrededor del H100, parece que Nvidia está usando viejas tácticas para proteger su foso defensivo.
    Por eso, más que profundizar demasiado en las particularidades del H100, habría que entender que “qué hardware quiere la IA” también incluye el contexto comercial.
  • AMD ya va por la segunda generación de su línea Versal.
    https://www.amd.com/en/products/accelerators/alveo/v80.html
    XDNA Architecture
    https://www.amd.com/en/technologies/xdna.html
  • ¿Google no lleva casi 10 años fabricando este tipo de dispositivos?

• Me llamó la atención esta parte: “Las mentiras de NVIDIA. Es una descripción tremendamente engañosa de la disposición real 128b swizzled wgmma. Como este diagrama me hizo perder tres semanas irrecuperables de mi vida, lo exhibo públicamente para avergonzarlos”.
  Me pregunto si a alguien le sorprenderá que una parte enorme del avance de la IA esté en ingeniería como la optimización de multiplicación de matrices, y que buena parte de esa ingeniería sea ingeniería inversa de chips NVIDIA.

  • La arquitectura en sí no hace una gran diferencia. Si entrenas un modelo lo suficientemente grande con datos lo suficientemente grandes, tiende a producir resultados parecidos sin importar la arquitectura. Por eso se puede decir que la mayor parte del avance de la IA se debe ahora a que podemos multiplicar matrices extremadamente rápido.

• Warp scheduler, cuatro cuadrantes, Tensor Memory Accelerator, disposición wgmma unswizzled…
  La frontera entre la terminología de GPU y el technobabble estilo Star Trek se vuelve cada vez más difusa.

  • Mientras leía el artículo, más o menos entendía de qué hablaban, pero decir “estamos haciendo warp de cuadrantes con un acelerador tensorial” realmente suena a Star Trek.
    También me ha pasado a veces al leer otros artículos. Me pregunto qué sentiría alguien que recibiera un enlace a un artículo de estos y lo leyera. Probablemente sería como entrar a una convención de fans de Trek discutiendo sobre el núcleo warp.
  • Ese comentario me hizo tomar distancia y mirar los términos con ojos nuevos, y me dio risa porque es totalmente cierto.

• Para reducir el consumo de energía de la inferencia de IA y aumentar la velocidad, parece que lo mejor sería pasar a circuitos analógicos aproximados.
  No se necesita una multiplicación y suma de punto flotante perfectas; solo hace falta un dispositivo que tome dos voltajes de entrada y entregue un voltaje de salida lo suficientemente cercano al resultado de la multiplicación.

  • Conozco a alguien que trabaja en esta dirección, y me dijo que los grandes obstáculos son cómo fabricar algo capaz de lógica analógica con las tecnologías existentes de fabricación de chips, diseñarlo para que no actúe como una antena, y la posibilidad de tener que ajustar finamente el modelo a ejecutar para cada chip, porque las tolerancias de fabricación varían de un chip físico a otro.
    La gran ventaja es que, en vez de representar un float16 con 16 líneas, se representa ese número con el voltaje de una sola línea. En teoría, incluso podría ser posible una precisión mucho mayor que float32. Además, como los valores pueden conectarse directamente sin cargarlos en una unidad aritmético-lógica, el ahorro potencial de área de die y de energía podría ser de varios órdenes de magnitud.
  • Creo que todavía falta mucho para que los circuitos analógicos sean útiles en la práctica, pero un lugar donde se podría aceptar la imprecisión son los circuitos digitales con ruido.
    Por ejemplo, aceptar que se invierta uno de cada millón de bits de salida y mejorar la relación rendimiento/energía. Sería difícil con float32, donde un solo valor infinito puede arruinarlo todo, pero con int8 parece tolerable que, cuando se esperaba un 0, de vez en cuando salga 128.
    [1] No estoy muy seguro de si las unidades de punto flotante para matrices del H100 realmente cumplen con IEEE 754.
  • Yendo un paso más allá, creo que necesitamos algo que se parezca a cómo funciona un cerebro biológico real, pero que también sea fácil de producir.
    Las redes neuronales biológicas no están ni cerca de estar completamente conectadas como las redes neuronales artificiales típicas; los coeficientes de conexión de entrada/salida de una neurona son menores que 10, así que son muy locales. En biología, hasta donde sabemos, no existe la retropropagación; en su lugar hay retroalimentación y recurrencia.
    También podría haber células o procesos auxiliares, aún desconocidos, que sean esenciales para funciones del sistema nervioso central. Incluso a alto nivel, probablemente exista una cantidad considerable de conectividad "hardcodeada", y parte de ella ya se conoce. Por ejemplo, las neuronas auditivas del oído están conectadas y ocurre algo parecido a una convolución para localizar la posición del sonido. Eso no es un fenómeno emergente, sino una función posible incluso sin entrenamiento.
    No es sorprendente, ya que la vida lo encontró a lo largo de miles de millones de años y un número comparable de generaciones. En teoría también sería posible hacerlo en software, pero considerando el billón o más de neuronas del cerebro de primates/humanos, sería extremadamente difícil incluso con las máquinas actuales de alrededor de mil núcleos. Incluso la "nube" no cumpliría con la conectividad y la latencia necesarias.
    Sería genial si con este enfoque se pudiera modelar con éxito algo del nivel de un gusano o un insecto.
  • Parece casi imposible satisfacer al mismo tiempo suficiente rango y precisión.
  • Sinceramente, parece una pesadilla para depurar.

• Este artículo me recordó lo mucho que disfruté la clase de programación paralela CS 149.

  • Kayvon y Kunle son increíbles. Tomé CS149 Parallel Programming hace dos semestres y fue excelente :)

• El estilo de escritura de este artículo me parece realmente impresionante, y tengo ganas de ver esto en la AMD MI300x. Si quieren usar tiempo en mi equipo, avísenme.

  • Me pregunto si has hecho mucho trabajo de IA con productos de AMD. No quiero gastar más de 2500 dólares en una RTX 4090, así que estoy considerando una RX 7900XTX para experimentar o para empezar.
    También me pregunto qué tan bien funcionaría en la práctica, si conviene ahorrar un poco más y comprar la XTX en lugar de la 7900 XT, y cuánto afectaría la usabilidad real tener menos VRAM.
  • Un buen texto debe ser claro y no ambiguo. Cuando uno habla, se puede interrumpir y pedir aclaración, pero un texto solo tiene una oportunidad de transmitir el mensaje.
    El lector no debería tener que ir a knowyourmeme.com para entender qué intentan decir los autores. Ni siquiera sé qué significa este título, y creo que por eso se aleja mucho del objetivo.
  • ¿En serio? Me trae PTSD de la época de Wallstreetbets.

• Me pregunto por dónde habría que empezar y qué hoja de ruta seguir para entender por completo un artículo como este.

  • Hay un buen curso para aprender programación en GPU. Para alrededor de la clase 4.0 ya puedes obtener los fundamentos necesarios: https://youtube.com/playlist?list=PLzn6LN6WhlN06hIOA_ge6Srgd...
    Y conviene escribir tú mismo un kernel CUDA que haga multiplicación vector-matriz. Si usas pycuda, puedes concentrarte en el kernel y escribir el resto en Python. Puedes decirle a ChatGPT que quieres crear tú mismo una implementación para multiplicar un vector de 4000 elementos por una matriz de 4000x12000, y pedirle que te guíe por todo el proceso.
    Para alquilar GPU, Runpod es una buena opción, y ahora tiene desde GPU de bajo costo hasta H100. Al principio basta con empezar con una GPU de gama baja.
  • Si quieres profundizar, quizá convenga ver la playlist de Spiral sobre multiplicación de matrices: https://www.youtube.com/playlist?list=PL04PGV4cTuIWT_NXvvZsn...
    Pasé 2 meses implementando y optimizando kernels de multiplicación de matrices con Spiral.

• El gráfico del README de GitHub (https://github.com/HazyResearch/ThunderKittens/blob/main/att...) marea demasiado. ¿Estos gráficos de barras onduladas son siquiera legales? :P

  • Estoy de acuerdo. Parece que intentaran meter alguna ilusión óptica. Preferiría ver solo los números, sin barras.
  • Parece el tema xkcd de matplotlib[1]. Aun así, coincido en que las ondas son demasiado exageradas.
    [1]: https://matplotlib.org/stable/gallery/showcase/xkcd.html#sph...

• ThunderKittens es un nombre excelente. Me gustaría ver que ThunderKittens aborde la retropropagación de FlashAttention, que es un orden de magnitud más difícil que la propagación hacia adelante.

  • Hay buenas noticias. TK ya incluye versiones causal y no causal optimizadas de la retropropagación de FlashAttention. Sería genial que les echaras un vistazo.
    causal: https://github.com/HazyResearch/ThunderKittens/blob/main/exa...
    no causal: https://github.com/HazyResearch/ThunderKittens/blob/main/exa...

• ¿Este tipo de investigación no es algo que ya vienen haciendo los equipos que fabrican NPUs hoy en día? Por ejemplo, el chip de Groq puede ofrecer el rendimiento que tiene porque usa una arquitectura dedicada a IA. Del lado de consumo, Apple Silicon también es bastante capaz.
  No soy de este campo, pero me parece que hay límites si solo se usan procesadores de propósito general que se comunican por rutas relativamente lentas. Repensar el diseño a nivel de hardware y, en última instancia, bajar los precios en el mercado de consumo parece una mejor estrategia a largo plazo.

  • No estoy tan seguro de que Apple Silicon sea bastante capaz del lado de consumo. Si miras el subreddit localllama de reddit, hay muchos posts de usuarios de CPU frustrados que intentan conseguir velocidades útiles.
    Cuando puedes comprar una GPU Nvidia por unos cientos de dólares, o una laptop gamer con 4050 y 6 GB de VRAM por 900 dólares, es difícil llamar “capaz” a la IA basada en CPU.
    En mi trabajo tampoco había GPU, así que probé con CPU, pero no era realista salvo usar modelos pequeños y esperar. Al final terminé pidiendo una computadora con GPU.
    “Técnicamente posible” y “realmente cómodo de usar” son cosas distintas. Nvidia fue realmente cómoda de usar; la CPU fue dolorosa y frustrante.