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

FlashAttention-4 es un rediseño del algoritmo de atención pensado específicamente para las GPU modernas, en las que los núcleos tensoriales mejoran su rendimiento mucho 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) depende de la velocidad de dos multiplicaciones de matrices: S = Q × K^T y O = P × V. Pero el análisis de la Blackwell B200 revela algo inesperado: el cuello de botella no está en los núcleos tensoriales, sino en las unidades de funciones especiales (SFU) que calculan la exponencial en el softmax (forward pass) y en el tráfico de la memoria compartida (backward pass). De Hopper (H100) a Blackwell (B200), el rendimiento de los núcleos tensoriales en BF16 pasó de 1 a 2.

25 petaflops, mientras que el número de SFU y el ancho de banda de la memoria compartida se mantuvieron igual. Esta asimetría rompe la optimización estándar: no se puede suponer sin más que los núcleos determinan por completo el rendimiento. En la práctica, se estorban entre sí.

Cómo resolvió el problema Together AI

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

• Nuevo pipelining: cauces por software que solapan al máximo el trabajo de los núcleos tensoriales, la SFU y la memoria, sin tiempos muertos.
• Emulación de la exponencial mediante polinomios: en lugar de la lenta SFU, el forward pass calcula la exponencial en las rápidas unidades FMA (fused multiply-add).
• TMEM y 2-CTA MMA: uso de la nueva memoria tensorial (256 KB por SM) más un modo en el que dos bloques de hilos trabajan sobre una misma operación matricial, lo que reduce el tráfico de la memoria compartida.

Las capacidades de Blackwell que lo hicieron posible

Blackwell incorporó varias funciones pensadas para este tipo de optimizaciones:

La memoria tensorial (TMEM) es un almacenamiento rápido en el chip (256 KB por cada SM), conectado directamente a los núcleos tensoriales. Los resultados intermedios pueden quedarse en la TMEM sin pasar por la más lenta memoria compartida, lo que reduce drásticamente la latencia de acceso.

Núcleos tensoriales asíncronos de 5.ª generación: cada operación la lanza un solo hilo y acumula los resultados en la TMEM. El tile máximo para BF16 es de 128×256×16 (aproximadamente el doble que en Hopper), lo que permite un pipelining más profundo sin desbordar los registros.

2-CTA MMA es un modo nuevo en el que dos bloques de hilos trabajan a la vez sobre una misma operación matricial. Reduce a la mitad la carga sobre la memoria compartida y disminuye el número de operaciones atómicas.

Cifras y resultados

En una Blackwell B200 con BF16, FlashAttention-4 alcanza 1605 TFLOPs/s (71 % de utilización). Es 1.3× más rápido que cuDNN 9.13 y 2.7× más rápido que Triton. Para situarlo: 1605 TFLOPs es casi la mitad del rendimiento máximo de Blackwell, pero se logra exprimiendo un kernel de atención complejo.

Qué significa

FlashAttention-4 muestra cómo trabajar en la era del escalado asimétrico de las GPU: en lugar de apretar tuercas en algoritmos antiguos, rehacerlos junto con las nuevas capacidades del hardware. El resultado parece un logro, pero es solo el comienzo de la adaptación a la nueva realidad del hardware.

*Meta está reconocida como organización extremista y prohibida en Rusia.