Paper #08 · Eficiencia e ingeniería
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
FlashAttention: atención exacta, rápida y eficiente en memoria, consciente de la E/S
- Autores:
- Dao, Fu, Ermon, Rudra y Ré
- Institución:
- Stanford University · University at Buffalo
- Año:
- 2022 · arXiv:2205.14135
En pocas palabras
FlashAttention calcula exactamente la misma atención que el Transformer original, pero reorganiza el cálculo en bloques para no escribir nunca la enorme matriz de atención en la memoria lenta de la GPU. Resultado: más velocidad, mucho menos memoria y contextos más largos.
Por qué importa hoy: Enseña que el cuello de botella no son solo los FLOPs: importan la HBM, la SRAM y el ancho de banda de memoria.
El problema
La atención crece de forma cuadrática con la longitud de la secuencia. Muchas propuestas aproximaban la atención para ahorrar operaciones, pero en la práctica no eran más rápidas: el tiempo se iba en mover datos entre la memoria grande y lenta (HBM) y la memoria pequeña y rápida dentro del chip (SRAM).
Ideas clave
Pensar en la E/S
El costo real de un algoritmo en GPU depende de cuántas veces lee y escribe en HBM, no solo de cuántas operaciones hace. La atención estándar escribe y vuelve a leer una matriz de tamaño N×N varias veces.
Tiling
Se parten Q, K y V en bloques que caben en SRAM y se procesan uno por uno, acumulando el resultado. La matriz completa de atención nunca se materializa en HBM.
Softmax en línea
Un truco numérico permite calcular el softmax por bloques, llevando el máximo y la suma acumulados, y corregir el resultado al llegar cada bloque nuevo.
Recalcular en vez de guardar
En la retropropagación, en lugar de guardar la matriz de atención, se recalcula por bloques. Se hacen más operaciones pero se leen muchos menos bytes, y el resultado es más rápido.
Resultados
- 15 % más rápido de extremo a extremo en BERT-large que el récord de MLPerf 1.1; 3 veces más rápido en GPT-2.
- Memoria lineal, en lugar de cuadrática, respecto a la longitud de la secuencia.
- Primeros Transformers en superar el azar en Path-X (16 mil tokens) y Path-256 (64 mil tokens).
Legado
FlashAttention y sus versiones 2 y 3 están integradas en PyTorch y en casi todos los motores de entrenamiento e inferencia. Son una razón directa de que hoy existan ventanas de contexto de cientos de miles de tokens.
Conceptos para leerlo
Explicación original en español elaborada por AI Atlas; no es una traducción del paper. El texto completo pertenece a sus autores.