La atención no era lenta por las cuentas
Durante años la explicación estándar de por qué la autoatención escala mal fue “es cuadrática”: con una secuencia de N tokens, calcular la atención cuesta del orden de N² operaciones. Es cierto, pero se queda corto. Cuando perfilas de verdad un Transformer entrenando, el cómputo casi nunca es lo que manda. La GPU pasa la mayor parte del tiempo esperando a la memoria.
Ese matiz es lo que cambió FlashAttention. No reduce el número de operaciones (sigue siendo O(N²) en cuentas), pero reordena cómo se hacen para que la GPU deje de ir constantemente a buscar datos a su memoria lenta. El resultado es una atención varias veces más rápida y con un consumo de memoria que pasa de crecer con N² a crecer de forma lineal. Y sin aproximar nada: la salida es exacta, equivalente token a token a la de la implementación ingenua.
La matriz que nunca deberías guardar
La atención estándar se calcula en tres pasos que todos hemos escrito alguna vez:
S = Q · Kᵀ— la matriz de puntuaciones, de tamaño N×N.P = softmax(S)por filas.O = P · V— la salida.
El problema vive en esa S. Para una secuencia de 8.000 tokens, la matriz N×N tiene 64 millones de entradas, y eso por cada cabeza y cada capa. La implementación ingenua la escribe entera en la memoria de la GPU, la vuelve a leer para el softmax, escribe P, la relee para multiplicar por V. Cada una de esas idas y vueltas cruza el cuello de botella.
Lo revelador es que casi todo ese trabajo de memoria es tráfico que podrías no hacer. La N×N solo existe como paso intermedio: lo que quieres es O, que es mucho más pequeña. Materializar S completa es guardar un resultado enorme para consumirlo dos líneas después y tirarlo.
Dos memorias: una rápida y otra lejos
Una GPU no tiene “la memoria” a secas, sino una jerarquía. Arriba está la SRAM on-chip: pequeñísima (decenas de KB por bloque de cómputo) pero con un ancho de banda descomunal. Abajo está la HBM, los decenas de GB que ves en nvidia-smi, mucho más lenta en comparación. La diferencia de ancho de banda entre ambas es de más de un orden de magnitud.
La atención ingenua trata la HBM como si fuera gratis. La matriz S de N×N no cabe en SRAM, así que vive en HBM, y con ella todo el tráfico de escribirla y releerla varias veces. El kernel queda memory-bound: las unidades de cómputo tienen las cuentas hechas y esperando; lo que falta es que lleguen los datos. Optimizar los FLOPs en ese escenario no sirve de nada, porque los FLOPs no son el problema.
| Atención ingenua | FlashAttention | |
|---|---|---|
| Matriz N×N | Se materializa entera en HBM | Nunca se materializa |
| Memoria extra | Crece con N² | Crece de forma lineal con N |
| Lecturas/escrituras en HBM | Varias pasadas sobre N×N | Una pasada sobre Q, K y V |
| Kernel | Varios kernels separados | Uno solo, fusionado |
| Resultado | Exacto | Exacto (idéntico) |
El truco: trocear y llevar la cuenta del softmax
FlashAttention nunca materializa S entera. Parte Q, K y V en bloques que sí caben en SRAM y recorre la atención por baldosas (tiling): carga un bloque de Q y un bloque de K/V, calcula ese trozo de puntuaciones dentro de la SRAM, lo acumula sobre la salida y lo descarta antes de pasar al siguiente bloque. La N×N completa jamás toca la HBM.
El obstáculo es el softmax. Normalizar una fila exige, en teoría, ver la fila entera: hay que restar el máximo (por estabilidad numérica) y dividir por la suma de los exponenciales. Si procesas K/V a trozos, nunca tienes la fila completa a la vez.
La solución es el softmax online: calcularlo de forma incremental. Arrastras dos estadísticos por fila —el máximo visto hasta ahora y la suma acumulada de exponenciales— y, cuando llega un bloque nuevo con un máximo mayor, reescalas el resultado parcial para que todo quede en la misma referencia. Al terminar el último bloque, el acumulado coincide exactamente con el softmax de haber tenido la fila entera de golpe.
# Softmax online sobre bloques de K/V (esquema mental de una fila de Q).
# La implementación real vive en un kernel CUDA fusionado; esto solo
# ilustra por qué no hace falta la matriz N×N completa.
import numpy as np
def flash_attention_row(q, k_blocks, v_blocks):
m = -np.inf # máximo visto hasta ahora
l = 0.0 # suma acumulada de exponenciales
o = np.zeros_like(v_blocks[0][0]) # salida parcial
for K, V in zip(k_blocks, v_blocks): # un bloque cada vez, cabe en SRAM
s = q @ K.T # puntuaciones solo de este bloque
m_new = max(m, s.max()) # ¿nuevo máximo global?
p = np.exp(s - m_new) # exponenciales estables
# reescala lo que ya llevábamos a la nueva referencia m_new
scale = np.exp(m - m_new)
l = l * scale + p.sum()
o = o * scale + p @ V
m = m_new
return o / l # normalización final exactaFíjate en que en ningún momento existe la matriz N×N. En cada iteración solo vive en SRAM el bloque de puntuaciones actual, que se consume y se olvida. La memoria extra ya no depende de N², sino de lo que ocupan los acumuladores.
Recomputar sale más barato que recordar
Queda el backward, el paso de entrenamiento. El gradiente necesita P, la matriz de probabilidades de atención. La implementación clásica la guardaba en la pasada hacia delante para reutilizarla al retropropagar; otra vez una N×N gigante viviendo en HBM.
FlashAttention hace lo contrario: no guarda P ni S. Se queda solo con los dos estadísticos del softmax (el máximo y la suma por fila, que son vectores, no matrices) y recomputa las puntuaciones que necesite a partir de Q, K y V durante el propio backward. Suena a herejía —¿repetir cuentas?— pero cuando estás memory-bound el trueque sale a favor: recomputar en SRAM cuesta menos que haber escrito la matriz en HBM y volver a leerla. Cambias FLOPs, que sobran, por tráfico de memoria, que es lo que escasea.
Qué ganas y qué no
Antes de reescribir nada, conviene tener claro dónde ayuda y dónde no:
- Lo que ganas. Entrenamiento e inferencia más rápidos en contextos largos, y sobre todo la posibilidad de usar contextos que antes no cabían en memoria, porque el coste en VRAM deja de ser cuadrático. Cuanto más larga la secuencia, mayor la ventaja.
- Sigue siendo O(N²) en cómputo. FlashAttention no convierte la atención en lineal; las cuentas siguen ahí. Lo que elimina es el tráfico de memoria innecesario. Para secuencias muy largas de verdad, eso se combina con atención dispersa o por ventanas, que sí atacan el coste de cómputo.
- No lo escribes tú. Vive en un kernel afinado a la arquitectura de la GPU. Ya está integrado en PyTorch (
scaled_dot_product_attention) y en los motores de inferencia serios. Tu trabajo es asegurarte de que la ruta rápida está activa, no reimplementar el softmax online. - Requiere la ruta soportada. Ciertas máscaras, sesgos o tipos de dato pueden hacer que la librería caiga a una implementación estándar sin avisar demasiado. Si esperabas la aceleración y no aparece, revisa que tu configuración encaje con lo que el kernel admite.
En resumen
La lección de FlashAttention trasciende la atención: el rendimiento en una GPU se decide muchas veces en el tráfico de memoria, no en el número de operaciones. La atención ingenua era lenta porque materializaba una matriz N×N enorme y la paseaba entre la memoria rápida y la lenta una y otra vez. FlashAttention trocea el problema para que quepa en SRAM, lleva la cuenta del softmax de forma incremental y recomputa en lugar de recordar. Mismas cuentas, mismo resultado, una fracción del tráfico.
La próxima vez que un kernel vaya lento y el nvidia-smi diga que la GPU está al 100 %, no des por hecho que faltan FLOPs. Mira cuánto de ese tiempo es cómputo real y cuánto es la tarjeta esperando datos. A menudo la palanca no está en hacer menos cuentas, sino en dejar de mover tantos bytes.
Fuentes
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — el paper original (Dao et al., 2022).