FlashAttention-4: cómo Together AI aceleró la atención en GPU Blackwell

FlashAttention-4 es un rediseño del algoritmo de atención especialmente para GPU modernas, donde los núcleos tensoriales crecen en rendimiento significativamente más rápido que la memoria y otros recursos.

Por qué el enfoque estándar ya no funciona

A primera vista, el rendimiento de la atención (attention) está controlado por la velocidad de dos multiplicaciones matriciales: S = Q × K^T y O = P × V. Pero el análisis de Blackwell B200 muestra algo inesperado: el cuello de botella no está en los núcleos tensoriales, sino en bloques funcionales especiales (SFU) para calcular la exponencial en softmax (forward pass) y en el tráfico de shared memory (backward pass).

De Hopper (H100) a Blackwell (B200), el rendimiento de los núcleos tensoriales BF16 creció de 1 a 2.25 petaflops, mientras que la cantidad de SFU y el ancho de banda de la shared memory permanecieron iguales. Esta asimetría rompe la optimización estándar —no se puede simplemente asumir que los núcleos determinan completamente el rendimiento. En realidad, se interfieren entre sí.

Cómo Together AI resolvió el problema

El equipo de investigación, junto con ingenieros de NVIDIA, Meta y Princeton, propuso tres ideas clave:

• Nuevo pipelineado —canalizaciones de software que superponen al máximo el trabajo de los núcleos tensoriales, SFU y memoria sin paros.

• Aproximación de exponencial mediante polinomios —en lugar del bloque SFU lento, el forward pass calcula la exponencial en unidades FMA rápidas (fused multiply-add).

• TMEM y 2-CTA MMA —uso de la nueva tensor memory (256 KB por SM) más un modo donde dos bloques de threads trabajan en una operación matricial, reduciendo el tráfico de shared memory.

Capacidades de Blackwell que lo hicieron posible

Blackwell agregó específicamente varias características para tales optimizaciones:

Tensor memory (TMEM) —almacenamiento rápido en el chip (256 KB por SM), conectado directamente a los núcleos tensoriales. Los resultados intermedios pueden permanecer en TMEM sin acceder a la shared memory más lenta, lo que reduce drásticamente la latencia de acceso.

Núcleos tensoriales asincronios de 5ª generación —cada núcleo se lanza con un único thread y acumula resultados en TMEM. El tile máximo para BF16 es 128×256×16 (aproximadamente 2 veces más grande que en Hopper), lo que permite un pipelineado más profundo sin desbordamiento de registros.

2-CTA MMA —un nuevo modo donde dos bloques de threads trabajan simultáneamente en una operación matricial. Esto reduce la carga en shared memory a la mitad y disminuye el número de operaciones atómicas.

Cifras y resultados

FlashAttention-4 en Blackwell B200 con BF16 logra 1605 TFLOPs/s (71% de utilización). Esto es 1.3× más rápido que cuDNN 9.13 y 2.7× más rápido que Triton. Para contextualizar: 1605 petaflops es casi la mitad del rendimiento máximo de Blackwell, pero se obtiene específicamente del complejo núcleo de atención.

Qué significa esto

FlashAttention-4 muestra cómo trabajar en la era del escalado asimétrico de GPU —no ajustar tuercas en algoritmos antiguos, sino rediseñarlos conjuntamente con las nuevas capacidades del hardware. El resultado parece un logro, pero es apenas el comienzo de la adaptación a la nueva realidad del hardware.

*Meta ha sido reconocida como una organización extremista y está prohibida en la RF.