Flash Attention

Flash Attention es un algoritmo optimizado para computar el mecanismo de atención en el núcleo de modelos transformers. Introducido por Tri Dao et al. de Stanford en un artículo de 2022, produce resultados numéricamente idénticos a la atención estándar pero reorganiza la computación para minimizar el movimiento de datos entre la memoria de alto ancho de banda (HBM) de GPU y la SRAM rápida en chip. Se clasifica como un algoritmo consciente de IO—su diseño está explícitamente motivado por el cuello de botella de ancho de banda de memoria del hardware moderno de GPU en lugar de reducir la complejidad aritmética de la operación de atención en sí.

La atención estándar materializa la matriz de atención completa N×N en HBM, donde N es la longitud de la secuencia, resultando en uso de memoria O(N²). Flash Attention usa una estrategia de tiling: procesa atención en bloques lo suficientemente pequeños para caber en SRAM, aplicando el algoritmo softmax en línea numéricamente estable para acumular resultados parciales incrementalmente. La matriz de atención completa nunca se escribe a HBM. Flash Attention 2 (2023) mejoró el paralelismo a través de la dimensión de secuencia y redujo operaciones de punto flotante no-matmul. Flash Attention 3 (2024) introdujo especialización de warp asíncrona y soporte de precisión FP8 dirigido a GPUs NVIDIA H100 y clase Hopper, impulsando el throughput más cerca de los picos teóricos del hardware.

La atención es el cuello de botella de memoria y computación dominante en entrenamiento e inferencia de transformers a longitudes de secuencia largas. Flash Attention reduce memoria de O(N²) a O(N) y entrega aceleraciones de pared de reloj 2–4x sobre implementaciones no optimizadas, traduciéndose directamente en costos de entrenamiento más bajos o más iteraciones experimentales por presupuesto de computación. Estas ganancias se componen a longitudes de secuencia más largas, haciendo que Flash Attention sea infraestructura esencial para modelos prácticos de contexto largo.

A partir de 2026, Flash Attention está integrado nativamente en PyTorch vía torch.nn.functional.scaled_dot_product_attention, Hugging Face Transformers, vLLM, y virtualmente todo marco principal de entrenamiento e inferencia de LLM. Es una dependencia de entrenamiento para LLaMA 3, Mistral, y la mayoría de otros modelos frontera. Ring attention, que distribuye tiling de Flash Attention a través de múltiples dispositivos vía comunicación all-to-all, extiende sus beneficios a secuencias que exceden memoria de GPU única, permitiendo las ventanas de contexto de millón-token ahora ofrecidas comercialmente.

Ejemplo

“Un equipo de investigación entrenando un transformer de 7-mil-millones-parámetros en secuencias de 32,000-tokens cambia de atención estándar de PyTorch a Flash Attention 2, reduciendo el uso de memoria GPU por paso en aproximadamente 60% y cortando el tiempo de entrenamiento por paso en 3x, permitiendo tamaños de lote más grandes y reduciendo a la mitad el tiempo de entrenamiento total de pared de reloj.