🤖 machine learning

Hardware-Aware FP4 FlashAttention-4

Este artículo presenta Hardware-Aware FP4 FlashAttention-4, un método que supera las limitaciones de los núcleos tensoriales FP4 de Blackwell mediante el empleo de Direct-P para la inferencia no causal y una ruta causal con gradientes FP8, logrando hasta un 2.13× de mayor rendimiento en el throughput de propagación hacia adelante y un 1.14× de mayor velocidad en las actualizaciones de entrenamiento en una sola GPU, evitando al mismo tiempo los problemas de divergencia observados en el entrenamiento MXFP4.

Autores originales: Robert Hu

Publicado 2026-09-04✓ Author reviewed ⓘ

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Robert Hu

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

En el mundo de la inteligencia artificial, los modelos más potentes dependen de un mecanismo llamado "atención" para comprender el contexto. Imagine a un lector escaneando un documento largo; la atención permite al modelo centrarse en las palabras más relevantes de una oración mientras ignora el resto, conectando ideas a través de vastas distancias. Para hacer esto, la computadora realiza una serie masiva de cálculos, comparando cada palabra con todas las demás para determinar sus relaciones. Durante años, estos cálculos se han realizado con números de alta precisión, similares a usar una regla con marcas diminutas y precisas para medir un edificio. Esto asegura que el modelo aprenda correctamente, pero es lento y consume enormes cantidades de energía. A medida que los modelos crecen, la necesidad de velocidad se vuelve crítica. Los investigadores han desarrollado recientemente chips capaces de utilizar números mucho más pequeños y toscos —valores de punto flotante de cuatro bits— para realizar estas comparaciones mucho más rápido. Sin embargo, simplemente cambiar a estos números más pequeños no es suficiente; el software que gestiona el flujo de datos también debe cambiar, o de lo contrario las ganancias de velocidad desaparecerán.

Un investigador de Graphcore ha abordado este cuello de botella específico con un nuevo método que llaman Direct-P. Su trabajo se centra en la pasada "hacia adelante" (forward pass) de la atención, donde el modelo procesa la información para generar una respuesta, y la pasada "hacia atrás" (backward pass), donde aprende de sus errores. El investigador descubrió que, si bien los nuevos chips pueden multiplicar números increíblemente rápido, el paso intermedio —convertir puntuaciones brutas en probabilidades— estaba ralentizando todo. Era como tener una línea de ensamblaje superrápida que se detenía constantemente porque un trabajador tenía que medir cuidadosamente cada pieza antes de pasarla. El investigador encontró que la forma estándar de manejar esta conversión, que implicaba un escalado y redondeo complejos, creaba un embotellamiento que anulaba los beneficios de velocidad del nuevo hardware.

Para resolver esto, el investigador rediseñó el proceso de conversión para que fuera directo y simplificado. En lugar de calcular una probabilidad precisa y luego redondearla, su método mapea las puntuaciones brutas directamente a los códigos específicos que utiliza el hardware. Esto elimina los pasos intermedios que causaban retrasos. Cuando probaron este enfoque en la última generación de chips para centros de datos, los resultados fueron sorprendentes. Para ciertas formas comunes de datos, el nuevo método procesó la información más del doble de rápido que el estándar anterior, que dependía de números de mayor precisión. Logró esto manteniendo la salida lo suficientemente precisa para tareas compleas como la generación de video y la comprensión del lenguaje. En pruebas con un modelo de generación de video, el nuevo método fue casi dos veces más rápido que el enfoque estándar, produciendo resultados que, aunque finitos y utilizables, mostraban una deriva medible y puntuaciones de similitud más bajas en comparación con la versión más lenta y precisa.

Sin embargo, la historia no trata solo de velocidad; también trata de lo que sucede cuando el modelo intenta aprender. El investigador exploró si podían usar estos números ultra rápidos y de baja precisión para todo el proceso de entrenamiento, incluyendo la pasada hacia atrás donde el modelo actualiza su conocimiento. Descubrieron que, si bien la pasada hacia adelante podía ejecutarse a máxima velocidad, la pasada hacia atrás requería un compromiso cuidadoso. Intentaron usar el formato de cuatro bits más rápido para los valores de probabilidad durante el entrenamiento, y el proceso de aprendizaje se volvió inestable y el modelo no logró mejorar. Los números se volvieron demasiado toscos, causando que la señal de aprendizaje se rompiera. Consecuentemente, el investigador determinó que para que el entrenamiento se mantuviera estable, debían usar un formato de ocho bits ligeramente más preciso para los valores de probabilidad, incluso si el resto del cálculo utiliza el formato de cuatro bits más rápido. Este enfoque híbrido les permitió acelerar todo el ciclo de entrenamiento aproximadamente un 14 por ciento en un solo chip potente, una ganancia significativa para modelos a gran escala.

El investigador también observó los límites físicos de los propios chips de computadora. Descubrió que la velocidad del cálculo ya no era el problema principal; el problema era cómo se almacenaban y movían los datos dentro del chip. El chip tiene un área de memoria pequeña y ultra rápida donde mantiene los números mientras trabaja en ellos. El investigador encontró que este espacio de memoria estaba completamente lleno, sin dejar espacio para solapar las diferentes etapas del cálculo. Era como una cocina donde la encimera está tan abarrotada de ingredientes que el chef no puede empezar a picar el siguiente vegetal hasta que termine con el actual, incluso si el cuchillo es increíblemente afilado. Debido a que el espacio de memoria estaba totalmente ocupado, el chip no podía trabajar en múltiples pasos a la vez, lo que limitaba qué tanto más rápido podría llegar a ser todo el proceso.

Este trabajo destaca un cambio crucial en cómo pensamos sobre cómo hacer que la inteligencia artificial sea más rápida. No basta con construir simplemente calculadoras más rápidas; todo el flujo de trabajo debe rediseñarse para coincidir con las capacidades del hardware. El investigador demostró que, al cambiar la forma en que se calculan las probabilidades y al gestionar cuidadosamente el flujo de datos, podían desbloquear mejoras masivas de velocidad. Sin embargo, también demostró que existen límites estrictos. Las restricciones de memoria del chip significan que, incluso con la matemática más rápida, hay un techo para cuánto solapamiento es posible. El camino a seguir, sugieren, no reside solo en una aritmética más rápida, sino en formas más inteligentes de organizar los datos para que los recursos del chip nunca se queden esperando. Este equilibrio entre la velocidad bruta y la gestión cuidadosa de los datos es lo que permitirá que la próxima generación de inteligencia artificial funcione de manera eficiente.

=== RESUMEN ===

Resumen Técnico: Hardware-Aware FP4 FlashAttention-4

Planteamiento del Problema

Las GPUs NVIDIA Blackwell introducen núcleos tensoriales de punto flotante de 4 bits (FP4) que aceleran significativamente la multiplicación de matrices. Sin embargo, en los mecanismos de atención, la aceleración suele verse limitada no por los productos de matrices ( $QK^T$ y $PV$ ), sino por la etapa intermedia de softmax. Esta etapa involucra la reducción de puntuaciones, la evaluación de exponenciales, la construcción de probabilidades y el escalado, lo cual debe completarse antes de que pueda comenzar el segundo producto de matrices.

Los intentos existentes de utilizar FP4 para todo el bloque de atención enfrentan dos obstáculos principales:

Temporización: La construcción del tensor de probabilidad ( $P$ ) a partir de las puntuaciones ( $S$ ) crea una dependencia de ruta crítica que retrasa el producto $PV$ .
Rango Numérico: Los formatos FP4 (específicamente NVFP4 y MXFP4) utilizan una carga útil de 4 bits (E2M1) con un rango dinámico limitado. Las probabilidades del softmax estándar abarcan un rango que es incómodo para esta carga útil, lo que conduce a subdesbordamientos (underflow), "escalas de cero" (donde bloques enteros desaparecen) o una pérdida significativa de precisión si no se gestionan cuidadosamente.

Además, en el entrenamiento causal, la pasada hacia atrás (backward pass) típicamente requiere reconstruir las probabilidades a partir de los estados guardados. Los enfoques que guardan puntuaciones de alta precisión o las recomputan en BF16 niegan las posibles aceleraciones del entrenamiento de baja precisión.

Metodología: Direct-P y Causal Backward

El artículo propone Direct-P, un método para permitir la inferencia directa en FP4 completo y una ruta de retropropagación causal cuantizada en el hardware Blackwell (GB200/B300).

1. Direct-P para Inferencia Directa

Direct-P modifica el intervalo entre la finalización del fragmento de puntuación $QK^T$ y la emisión de la instrucción $PV$ . Retiene el actual pipeline de dos consultas de HAO AI Lab y el protocolo de propiedad de Tensor Memory (TMEM), pero optimiza la generación de probabilidades:

Mapeo Directo: En lugar de calcular una exponencial de alta precisión y luego redondear, Direct-P trata la construcción de la probabilidad como un problema de clasificación de código. Mapea las puntuaciones logarítmicas normalizadas directamente a las 8 magnitudes E2M1 no negativas mediante un clasificador afín (o exponenciales nativas selectivas). Esto acorta la ruta crítica.
Normalización Consistente: El denominador (normalizador) para el softmax se acumula directamente a partir de los códigos de probabilidad FP4 redondeados y las escalas de bloque que serán consumidos por la instrucción $PV$ . Esto asegura que el numerador y el denominador describan el mismo operador aproximado, evitando discrepancias entre sumas de alta precisión y cargas útiles de baja precisión.
Protección de Rango (Range Guarding): Para capas con logits extremos (ej. >500), un mecanismo de anclaje muestreado evita subdesbordamientos subnormales sin requerir un re-escaneo completo de las puntuaciones.
Selección de Formato: El método utiliza NVFP4 para $Q$ y $K$ (debido al escalado fino) y MXFP4 para $P$ y $V$ (debido a sus escalas de bloque de potencia de dos que se alinean bien con los fragmentos de productor $N_{32}$ ).

2. Causal Backward Cuantizado

Para el entrenamiento, el artículo introduce una ruta donde la pasada hacia adelante guarda la carga útil cuantizada de $Q/K$ , las escalas y el normalizador log-sum-exp (LSE).

Reconstrucción: La pasada hacia atrás reconstruye las probabilidades directamente desde estas operandos cuantificados guardados en lugar de crear una ruta de puntuación BF16 separada.
Formatos de Gradiente: Los epílogos de proyección publican vistas orientadas a filas/columnas de FP8 de $Q, K, V$ para alimentar los productos de matrices de gradiente.
Restricción de Estabilidad: Aunque la ruta hacia adelante utiliza MXFP4 para $P/V$ , los experimentos de entrenamiento distribuido revelaron que las trayectorias de MXFP4 P/V divergen. Por consiguiente, la ruta de entrenamiento retenida utiliza FP8 para $P$ y $V$ en la pasada hacia atrás para asegurar la estabilidad, manteniendo $Q/K$ en NVFP4.

Contribuciones Clave

Implementación de Direct-P: Un método que mapea puntuaciones directamente a códigos de probabilidad MXFP4 y normaliza utilizando los mismos valores representados, logrando hasta un 2.13× de throughput de inferencia forward en BF16 en GB200 para formas favorables, con una aceleración de media geométrica de 2.02× en la rejilla D128 probada.
Ruta de Entrenamiento Causal Cuantizada: Una pasada hacia atrás que reconstruye las probabilidades directamente de los estados de cuantificación forward guardados, acelerando una actualización completa de un modelo de 8 billones de parámetros en una sola GPU hasta en 1.14×.
Perfilado de Hardware: Identificación de la propiedad de Tensor Memory (TMEM) como el límite principal de solapamiento (overlap). El estudio demuestra que tener núcleos tensoriales más rápidos es insuficiente si el diseño de almacenamiento (512 columnas en D128) impide que la siguiente operación $QK$ se solape con la fase actual de $PV$ .

Resultados

Rendimiento de Inferencia Forward

Throughput: En GB200, la ruta "rápida" de Direct-P logra una aceleración de media geométrica de 2.02× sobre la línea base de BF16 de HAO, alcanzando un pico de 2998 TFLOP/s (S32768/H24).
Precisión: Direct-P intercambia algo de precisión por velocidad. En pruebas de softmax Gaussiano, la ruta "rápida" tiene una similitud de coseno media de ~0.944 contra BF16, comparado con ~0.990 para la ruta NV/FP8 de HAO.
Impacto en el Modelo:
- ViT/BERT: En evaluaciones de entrada fija, la ruta "rápida" mantiene la precisión de la tarea (ej. la precisión top-1 de ViT S4096 coincide con BF16 al 88.5%) a pesar de la menor similitud de coseno a nivel de operador.
- Wan Video Diffusion: El método permanece finito y funcional para 20 pasos de difusión. Sin embargo, el error se acumula a través de los pasos; el modelo de 14B en 20 pasos muestra un L2 relativo de ~0.53 comparado con BF16, indicando deriva pero una generación funcional.

Rendimiento de Entrenamiento

Single-GPU: La ruta cuantizada (proyecciones NVFP4 + P/V en FP8) acelera la subcapa de atención completa en 1.25× y la actualización completa del modelo de 8B en 1.14× con un tamaño de lote de 4.
Entrenamiento Distribuido: En un experimento emparejado de 64 GPUs sobre 100 mil millones de tokens, la ruta P/V en FP8 logró una mejora de throughput de 1.112× sobre BF16.
Estabilidad: El artículo nota explícamente que las trayectorias de MXFP4 P/V divergieron en entornos distribuidos. La ruta de entrenamiento exitosa depende de usar FP8 para el producto $P/V$ , lo que significa que la ruta de entrenamiento no es "pura" FP4 de extremo a extremo.

Significancia y Reivindicaciones

El artículo afirma que Direct-P logra hacer práctica la atención FP4 completa en el forward al acortar la ruta crítica entre los productos de matrices. Demuestra que, si bien los núcleos de matriz FP4 son más rápidos, el cuello de botella se ha desplazado al movimiento de datos, la construcción de la probabilidad y la propiedad de la memoria.

Puntos Clave:

Límites de Hardware: La barrera principal para mayores aceleraciones no es el rendimiento aritmético, sino la propiedad de TMEM. El diseño actual de D128 consume las 512 columnas de TMEM, impidiendo que la siguiente operación $QK$ se solape con la cola de la fase $PV$ .
Paradas de Disponibilidad (Readiness Stalls): Una baja actividad de los núcleos tensoriales es a menudo un síntoma de la disponibilidad de los operandos (esperando la construcción de la probabilidad) más que de la falta de capacidad de cómputo.
Compromisos de Entrenamiento: Si bien la inferencia forward en FP4 es viable, el entrenamiento causal estable actualmente requiere una concesión hacia FP8 para los operandos de probabilidad/valor para prevenir la divergencia. El artículo sugiere que futuros formatos como Unsigned E5M3 (UE5M3) podrían permitir un entrenamiento totalmente en FP4, pero esto no se evalúa aquí.

El trabajo concluye que la aritmética de matriz más rápida por sí sola es insuficiente; lograr una mayor eficiencia requiere cambios arquitectónicos que permitan ventanas de solapamiento más grandes (ej. destinos de puntuación adicionales) junto con algoritmos de baja precisión optimizados.