Un ingeniero publica una guía detallada para construir, partiendo de cero, un kernel denso de atención para la GPU Nvidia B200 (arquitectura Blackwell) en CUDA con algo de PTX. El recorrido, ilustrado con 60 diagramas, avanza a través de 14 kernels progresivos hasta alcanzar el 94,4 % del rendimiento de FlashAttention-4 en formas de 4K, 8K y 16K empleadas en el propio paper de FA4.
El texto arranca con un kernel de referencia que produce resultados correctos pero no optimizados, y va incorporando una mejora por capítulo. La motivación central es la asimetría del hardware Blackwell: el throughput de los tensor cores se duplicó respecto a la generación anterior, mientras que las unidades de exponencial (exp) apenas cambiaron. Esto provoca que el softmax, pese a realizar pocos FLOPs, consuma tantos ciclos como las multiplicaciones de matrices en los tamaños de tile elegidos, y se convierta en el cuello de botella principal.
La estrategia de optimización combina dos vías: abaratar el trabajo del softmax y solaparlo con los multiplicadores matriciales (MMA). El trabajo se paraleliza asignando un tile de salida a cada CTA ( Cooperative Thread Array), que carga su tile de Q una sola vez e itera secuencialmente sobre los pares de tiles de K y V, manteniendo un acumulador privado del resultado.
La guía incluye un apartado introductorio opcional para principiantes en CUDA, no exige disponer de una B200 física y culmina con un proyecto capstone en el que el kernel final se integra en un modelo de generación de vídeo. Todo el código está disponible en el repositorio b200-attention de GitHub del autor.