Atención exacta, acelerada al no escribir nunca la enorme cuadrícula. · Source

Las mismas cuentas de la atención — solo que nunca se anotan.

El verdadero costo de la atención no es multiplicar — es arrastrar una tabla gigante de puntajes hacia y desde la memoria lenta. Flash attention no cambia nada de las cuentas y todo sobre dónde ocurren. La misma respuesta, en una fracción del tiempo. Como resolver un cubo a toda velocidad: el mismo rompecabezas, los mismos movimientos — la ventaja es puro método.

El cuello de botella no son las cuentas. Es el traslado.

La atención estándar arma la tabla completa de cada palabra puntuada contra todas las demás — luego escribe esa cuadrícula gigante en la memoria lenta, la vuelve a leer para normalizar y la lee otra vez para ponderar los valores. Tres viajes de ida y vuelta con el objeto más grande del modelo, mientras el chip termina las cuentas y solo espera. Como una alfombra desenrollada: demasiado grande para la sala, así que vive en el pasillo — y cada vistazo significa otro viaje.

Un chip veloz, ahogado por una tubería lenta.

T \approx \max\!\left(\dfrac{\text{FLOPs}}{\text{compute rate}},\ \dfrac{\text{bytes moved}}{\text{bandwidth}}\right)

Un procesador hace aritmética mucho más rápido de lo que puede traer los números que masticar. Así que el tiempo total es lo que sea más lento: hacer las cuentas o mover los bytes. La atención está limitada por la memoria — las cuentas están listas, los datos no — así que recortar viajes, no recortar cuentas, es lo que la acelera. Como una bomba potente en una manguera fina: podría mover un océano, pero la tubería decide el caudal.

Así que no lo cargues todo. Trabaja un bloque a la vez.

S_{ij} = \dfrac{Q_i K_j^{\top}}{\sqrt{d}}, \qquad Q_i,\, K_j \in \mathbb{R}^{B\times d}

La solución empieza simple: nunca acerques toda la cuadrícula. Corta las Keys y los Values en bloques lo bastante pequeños para caber en el diminuto bloc rápido del chip. Trae un bloque, haz su parte del trabajo, suéltalo, busca el siguiente. En palabras simples: cada bloque de B filas cabe en la memoria rápida, así que manejas solo un par de bloques a la vez. Como acarrear leña: no arrastras toda la pila hasta el hogar — traes una brazada.

Pero softmax necesita toda la fila a la vez.

\mathrm{softmax}(x)_i = \dfrac{e^{\,x_i - m}}{\sum_{j=1}^{n} e^{\,x_j - m}}, \qquad m = \max_{j} x_j

Aquí está el problema. Para convertir puntajes en pesos, softmax divide cada uno por el total sobre cada key — y por seguridad primero resta el puntaje más alto de la fila, para que las exponenciales nunca se disparen. Tanto el máximo como la suma necesitan toda la fila. Pero solo tienes un bloque. ¿Cómo normalizas contra números que no has visto? Como repartir partes justas de una sola jarra: no puedes llenar el primer vaso hasta saber cuántos hay.

El truco: lleva un máximo móvil y reajusta sobre la marcha.

\begin{aligned} m_{\text{new}} &= \max(m_{\text{old}},\, \tilde m) \\ \ell_{\text{new}} &= e^{\,m_{\text{old}} - m_{\text{new}}}\,\ell_{\text{old}} + \sum_{i} e^{\,x_i - m_{\text{new}}} \end{aligned}

No necesitas toda la fila primero — necesitas un conteo móvil. Lleva el puntaje más alto visto hasta ahora y el total acumulado. Cuando un bloque nuevo trae un puntaje mayor, pon el viejo total en la nueva escala con un solo factor de reducción, y luego suma el bloque; la salida se reescala igual. El resultado es idéntico a hacerlo todo de una vez. Como la barra de salto de altura: súbela cuando llega un mejor saltador, y vuelve a medir cada marca anterior contra la nueva altura.

La misma respuesta. La cuadrícula gigante nunca existió.

\text{memory: } O(n^2)\ \longrightarrow\ O(n), \qquad \mathbf{o}_{\text{flash}} = \mathbf{o}_{\text{exact}}

Júntalo todo y la tabla n×n nunca se almacena — solo un máximo móvil, una suma móvil y la salida, todo pequeño. La memoria cae de cuadrática a lineal, y la aceleración viene de muchísimos menos viajes a la memoria lenta, no de recortar cuentas. Y esto no es una aproximación: el resultado es atención exacta. Como plegar una grulla de origami: una hoja, la forma terminada, sin recortes sobre la mesa.

🌱 Quizá nunca fue el pensar lo que era lento.

Flash attention no agrega ninguna idea nueva — solo respeta dónde vive el cálculo, reordenando los mismos pensamientos para el cuerpo que los aloja. Entonces, ¿cuánto de la velocidad de una mente es el pensar en sí — y cuánto es solo qué tan rápido alcanza lo que ya sabe?