FlashAttention-3 (ES)

FlashAttention-3 es un algoritmo para optimizar el mecanismo de atención (attention) en las redes neuronales transformadoras, diseñado para aprovechar al máximo las capacidades de hardware de la arquitectura de GPU NVIDIA Hopper (H100). El algoritmo fue presentado en 2024 por un grupo de investigadores de Colfax Research, Meta, NVIDIA, Georgia Tech, la Universidad de Princeton y Together AI. El trabajo fue aceptado en la conferencia NeurIPS 2024 y destacado como spotlight.

FlashAttention-3 es la tercera iteración en la familia de algoritmos, siguiendo a FlashAttention (2022) y FlashAttention-2 (2023). Su objetivo principal es acelerar significativamente el entrenamiento y la inferencia de los grandes modelos de lenguaje (LLM), manteniendo al mismo tiempo la precisión de los cálculos.

Introducción y antecedentes

El problema del mecanismo de atención

El componente clave de los transformadores es el mecanismo de autoatención (self-attention), pero su complejidad computacional y su consumo de memoria crecen de forma cuadrática (O(n²)) con el aumento de la longitud de la secuencia de entrada (n). Esto crea un importante "cuello de botella", ya que las GPU modernas están optimizadas para multiplicaciones de matrices rápidas, pero el cálculo de funciones exponenciales (como en Softmax) es órdenes de magnitud más lento. Además, en una implementación ingenua, se debe almacenar un gran tensor de atención intermedio en la memoria de la GPU, lo que limita la escalabilidad de los modelos.

FlashAttention y FlashAttention-2

Para resolver este problema, en 2022 se propuso FlashAttention, que redujo el número de accesos a la lenta memoria global (HBM) mediante dos técnicas:

• Procesamiento en bloques (tiling): Los cálculos se dividen en bloques (tiles) que se procesan en la memoria rápida en el chip (SRAM).
• Fusión de operaciones: Todas las operaciones (multiplicación de matrices, Softmax) se realizan en un solo núcleo de la GPU sin escribir resultados intermedios en la memoria global.

Esto permitió reducir la complejidad de la memoria de cuadrática a lineal y aceleró los cálculos de 2 a 4 veces.

En 2023, se introdujo una versión mejorada, FlashAttention-2, que optimizó la paralelización de los cálculos. En las GPU de la arquitectura NVIDIA Ampere (A100), alcanzó aproximadamente el 70% del rendimiento teórico máximo. Sin embargo, en la arquitectura más reciente NVIDIA Hopper (H100), su eficiencia fue significativamente menor, de alrededor del 35%. Esto se debió a que el algoritmo no aprovechaba las nuevas capacidades de hardware de Hopper, lo que impulsó la creación de FlashAttention-3.

Nuevas capacidades de hardware de la GPU Hopper (H100)

La arquitectura NVIDIA Hopper introdujo varias características nuevas que FlashAttention-3 utiliza para lograr el máximo rendimiento:

• WGMMA (Warpgroup Matrix Multiply-Accumulate): Un nuevo tipo de instrucción para los núcleos tensoriales que realiza multiplicaciones de matrices con casi el doble de rendimiento en comparación con la arquitectura Ampere.
• TMA (Tensor Memory Accelerator): Un módulo de hardware que acelera la transferencia de datos entre la memoria global (HBM) y la memoria compartida (shared memory). El TMA realiza cálculos de direcciones automáticamente, liberando los núcleos de cómputo.
• Formato FP8: Soporte de hardware para el formato de datos de punto flotante de 8 bits, que duplica el rendimiento teórico en comparación con FP16, pero conlleva el riesgo de pérdida de precisión debido a su rango dinámico limitado.

Innovaciones técnicas de FlashAttention-3

El algoritmo implementa tres métodos clave de optimización, diseñados específicamente para la arquitectura Hopper:

1. Ejecución asíncrona y especialización de warps

FlashAttention-3 utiliza el principio de warp-specialization, donde diferentes grupos de hilos (warps) en la GPU se especializan en diferentes tareas:

• Warps productores (Producer warps): Cargan datos desde la memoria global utilizando el TMA.
• Warps consumidores (Consumer warps): Realizan multiplicaciones de matrices en los núcleos tensoriales.

Gracias a la asincronía de hardware de Hopper, estas operaciones se superponen en el tiempo. Mientras un grupo de warps realiza cálculos, otro carga datos para el siguiente bloque en paralelo. Este enfoque de canalización (pipeline), organizado según el principio de "ping-pong" (ping-pong scheduling), permite ocultar las latencias de operaciones lentas (como Softmax) y cargar al máximo todos los módulos funcionales de la GPU.

2. Minimización de las operaciones de memoria

El algoritmo mantiene la ideología de tiling de las versiones anteriores, pero utiliza activamente el TMA para cargar de forma asíncrona los siguientes bloques de datos en paralelo con los cálculos actuales. La transferencia de datos desde la lenta memoria HBM a la rápida SRAM se realiza efectivamente "a la sombra" de los cálculos principales, lo que reduce el tiempo que la GPU pasa inactiva esperando datos.

3. Baja precisión (FP8) con reducción del error de cuantización

El cambio a FP8 duplica la velocidad, pero puede llevar a una pérdida significativa de precisión debido a la cuantización. Para combatir esto, los desarrolladores implementaron el método de procesamiento incoherente (incoherent processing). Su esencia es la siguiente:

1. Antes de calcular la atención, los vectores de características (consultas Q y claves K) se multiplican por una matriz ortogonal aleatoria (por ejemplo, una matriz de Hadamard).
2. Esta transformación "dispersa" los valores con una magnitud anormalmente grande (valores atípicos) por todas las coordenadas, ecualizando su distribución.
3. Después de esto, se realiza la cuantización a FP8, que ahora se produce con un error menor.
4. Dado que la transformación es ortogonal, no distorsiona el resultado final de la atención (QKᵀ), ya que el efecto de la matriz se anula durante la multiplicación.

Esta técnica permitió reducir el error de cálculo de la atención en FP8 aproximadamente 2,6 veces en comparación con la aplicación estándar de FP8 sin transformaciones.

Rendimiento e importancia

La aplicación de estas técnicas permitió a FlashAttention-3 lograr una superioridad significativa sobre las versiones anteriores en la GPU H100:

• Aceleración de 1,5 a 2 veces en comparación con FlashAttention-2.
• Alta utilización de la GPU: Alcanza entre el 75% y el 85% del rendimiento máximo teórico de la H100.
• Rendimiento (throughput):
  • Hasta 740–840 TFLOPS para precisión media (FP16/BF16).
  • Hasta 1,2–1,3 PFLOPS (petaflops) utilizando precisión de 8 bits (FP8).

La alta eficiencia de FlashAttention-3 impacta directamente en el desarrollo y la aplicación de los LLM:

• Reducción del tiempo de entrenamiento: Una aceleración de la atención del 75-100% reduce significativamente el tiempo de entrenamiento de los modelos, que puede durar semanas o meses.
• Ampliación de la ventana de contexto: Los modelos pueden procesar eficientemente secuencias más largas (cientos de miles de tokens), lo cual es crucial para analizar documentos extensos o código.
• Uso racional de los recursos: Permite alcanzar el mismo rendimiento con menos GPUs o lograr mayor velocidad con el mismo hardware, lo que reduce el costo de despliegue de los modelos.

Disponibilidad e integración

Los autores publicaron el código fuente de FlashAttention-3 bajo una licencia de código abierto en GitHub. Se espera su integración en los principales marcos de aprendizaje profundo, como PyTorch y las bibliotecas de Hugging Face Transformers, lo que hará que la tecnología sea accesible para una amplia gama de desarrolladores e investigadores. Las versiones anteriores ya se han convertido en un estándar de facto en la industria, y es probable que FlashAttention-3 continúe esta tendencia.

Enlaces

• Repositorio oficial de FlashAttention en GitHub
• Blog de Together AI con el anuncio de FlashAttention-3

Bibliografía

• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. arXiv:2410.01359.
• Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.