FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Guía de lectura

Problema y motivación

Los Transformers tienen una complejidad temporal y de memoria cuadrática respecto a la longitud de la secuencia $N$. Aunque los métodos de atención aproximada reducen el recuento teórico de FLOPs, a menudo no logran mejoras de velocidad en tiempo real porque las operaciones en las GPUs modernas están limitadas principalmente por las transferencias de memoria (E/S) en lugar de las operaciones aritméticas.

En la atención estándar, el cálculo de $\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }$, $\mathbf{P}=\mathrm{softmax}(\mathbf{S})$ y $\mathbf{O}=\mathbf{PV}$ escribe matrices intermedias de tamaño $N\times N$ en la lenta memoria de alto ancho de banda (HBM) y las vuelve a leer múltiples veces. Esto genera $\Theta (Nd+N^2)$ accesos a la HBM y una huella de memoria de $O(N^2)$, lo que restringe el tamaño de las ventanas de contexto prácticas.

Ideas clave

• Conciencia de la E/S (IO-Awareness): Considerar explícitamente la jerarquía de memoria de la GPU: la memoria SRAM en el chip, rápida y pequeña (por ejemplo, 192 KB/SM con un ancho de banda de ~19 TB/s en A100), frente a la memoria HBM, grande y lenta (40–80 GB con un ancho de banda de 1.5–2.0 TB/s).
• Particionado o Tiling (Cálculo de softmax por bloques): Dividir las entradas $\mathbf{Q},\mathbf{K},\mathbf{V}$ en bloques que quepan por completo en la SRAM. Al registrar las estadísticas acumuladas del máximo por fila y del normalizador de la softmax, la reducción de softmax se calcula incrementalmente a través de los bloques sin materializar la matriz completa de $N\times N$ en la HBM.
• Recomputación en el pase hacia atrás: Evitar guardar la matriz de atención intermedia de $N\times N$ para la retropropagación. En su lugar, almacenar únicamente la salida $\mathbf{O}$ y las estadísticas de softmax $(\ell ,m)$, recomputando $\mathbf{S}$ y $\mathbf{P}$ en el chip bloque a bloque en la SRAM durante el pase hacia atrás.
• Fusión de kernels: Fusionar las multiplicaciones matriciales, el enmascaramiento, la softmax y el dropout en un solo kernel de CUDA para eliminar transferencias redundantes de ida y vuelta a la HBM.

Método (con las ecuaciones clave)

Dados $\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$, FlashAttention establece los tamaños de bloque $B_c=\lceil M/(4d)\rceil$ y $B_r=\min (\lceil M/(4d)\rceil ,d)$, donde $M$ representa la capacidad de la SRAM.

Seguimiento incremental de la Softmax

Para un vector descompuesto en subvectores $x^{(1)}$ y $x^{(2)}$, las estadísticas acumuladas se actualizan mediante: $$m(x)=\max (m(x^{(1)}),m(x^{(2)})),$$ $$\ell (x)=e^{m(x^{(1)})-m(x)}\ell (x^{(1)})+e^{m(x^{(2)})-m(x)}\ell (x^{(2)}).$$

Para el bloque $j$ de claves/valores y el bloque $i$ de consultas:

1. Calcular las puntuaciones no normalizadas en SRAM: ${\mathbf{S}}_{ij}={\mathbf{Q}}_i{\mathbf{K}}_j^{\top }\in {\mathbb{R}}^{B_r\times B_c}$.
2. Calcular las estadísticas locales del bloque: ${\tilde{m}}_{ij}=\mathrm{rowmax}({\mathbf{S}}_{ij})$, ${\tilde{\mathbf{P}}}_{ij}=\exp ({\mathbf{S}}_{ij}-{\tilde{m}}_{ij})$, y ${\tilde{\ell }}_{ij}=\mathrm{rowsum}({\tilde{\mathbf{P}}}_{ij})$.
3. Actualizar los normalizadores globales acumulados: $$m_i^{\mathrm{new}}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{\mathrm{new}}=e^{m_i-m_i^{\mathrm{new}}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\ell }}_{ij}.$$
4. Actualizar el bloque de salida acumulado ${\mathbf{O}}_i\in {\mathbb{R}}^{B_r\times d}$: $${\mathbf{O}}_i\leftarrow \mathrm{diag}({\ell }_i^{\mathrm{new}}{)}^{-1}\left(\mathrm{diag}({\ell }_i)e^{m_i-m_i^{\mathrm{new}}}{\mathbf{O}}_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\mathbf{P}}}_{ij}{\mathbf{V}}_j\right).$$

Análisis de complejidad

• Complejidad de E/S: La atención estándar requiere $\Theta (Nd+N^2)$ accesos a HBM. FlashAttention requiere $\Theta (N^2{d}^2{M}^{-1})$ accesos a HBM. Los autores demuestran que ningún algoritmo de atención exacta puede lograr asintóticamente $o(N^2{d}^2{M}^{-1})$ accesos para todo $M\in [d,Nd]$.
• FLOPs y Memoria: FlashAttention realiza $O(N^{2}d)$ FLOPs (con una cantidad ligeramente superior de FLOPs en el pase hacia atrás debido a la recomputación) y solo requiere $O(N)$ de memoria auxiliar más allá de las entradas y salidas.

Extensión: Block-Sparse FlashAttention

Restringir la evaluación a bloques no nulos según una máscara de dispersión predefinida $\mathbf{M}\in \{0,1{\}}^{N/B_r\times N/B_c}$ logra una complejidad de E/S de $\Theta (Nd+N^2{d}^2{M}^{-1}s)$, donde $s$ es la fracción de bloques no nulos.

Resultados (con las cifras más importantes)

• Aceleración del entrenamiento:
  • BERT-large (longitud de secuencia 512): Supera el récord de velocidad de entrenamiento de MLPerf 1.1 en un 15% (17.4 min frente a 20.0 min en $8\times$ A100s).
  • GPT-2 (longitud de secuencia 1K): Logra una aceleración de hasta 3.5x respecto a HuggingFace y de hasta 1.8x respecto a Megatron-LM.
  • Long-Range Arena (LRA): Ofrece una aceleración de 2.4x en comparación con la atención estándar; Block-sparse FlashAttention alcanza una aceleración de 2.8x.
• Capacidades para contextos más largos:
  • GPT-2 con contexto de 4K: Se entrena más rápido que Megatron-LM con contexto de 1K mientras mejora la perplejidad de validación en 0.7 puntos (17.5 frente a 18.2).
  • Clasificación de documentos largos: Escalar el contexto de 512 a 16K proporciona una mejora de 4.3 micro-$F_1$ en MIMIC-III (de 52.8 a 57.1) y de hasta 8.5 puntos en ECtHR (de 72.2 a 80.7 en 8K).
  • Path-X y Path-256: FlashAttention es el primer Transformer capaz de resolver Path-X (longitud de secuencia de 16K) con un 61.4% de precisión. Block-sparse FlashAttention logra un 63.1% de precisión en Path-256 (longitud de secuencia de 64K).

Limitaciones y preguntas abiertas

• Sobrecarga de ingeniería en CUDA: Implementar kernels conscientes de la E/S requiere programación en CUDA de bajo nivel y a medida, en lugar de utilizar abstracciones estándar de PyTorch o compiladores, lo que dificulta la portabilidad entre diferentes arquitecturas de hardware.
• Atención multi-GPU: El esquema óptimo de E/S para un único nodo no se extiende de manera automática a la atención distribuida multi-GPU a través de interconexiones (NVLink, InfiniBand).
• Cota inferior de la atención exacta: Aunque es óptimo para valores generales de $M$, demostrar cotas inferiores de complejidad parametrizada más detalladas en función de $M$ sigue siendo un problema abierto.