La parte 3 acababa en un cuadrado. En un bloque de decodificador hay dos tensores con T dos veces en la forma, [B, H, T, T], es decir, un cuadrado T × T por cabeza, con H cabezas y B secuencias. Con los valores por defecto del gráfico de aquella parte, suponían el 89 % de la memoria del bloque. En esta leo lo que los llena, que es una media ponderada. La media en sí no tiene nada de especial. Lo nuevo es que los pesos se calculan a partir de la entrada.
He escrito la atención en cuatro líneas y la he contrastado con F.scaled_dot_product_attention, la implementación del propio torch. Luego he leído LlamaAttention.forward con las formas anotadas al margen y he calculado cuánto cuesta en Llama 3.1 8B con las longitudes de contexto que se prometen.
Un diccionario que coincide en parte
Buscar en un diccionario es una operación exacta: la clave coincide o no coincide, y obtienes un único valor. La atención hace lo mismo con una coincidencia difusa. Todas las claves coinciden un poco y obtienes la media de todos los valores, ponderada por lo bien que ha coincidido cada clave.
Esas ponderaciones son los pesos de atención. No son parámetros ni se guardan en ningún sitio, porque el modelo los calcula a partir de la entrada en cada llamada. Hasta llegar a la sección del coste, cuando diga «pesos» me referiré a estos. Tres matrices aprendidas convierten cada fila de la entrada en una query, una key y un value. La query se compara con todas las keys y las puntuaciones resultantes se convierten en pesos:
q, k, v = x @ Wq, x @ Wk, x @ Wv # [B, T, d] eachscores = q @ k.transpose(-1, -2) / sqrt(d) # [B, T, T]weights = scores.softmax(-1) # every row sums to 1out = weights @ v # [B, T, d]Fíjate en las formas. x es el [B, T, d] de la parte 3 y las tres proyecciones lo conservan. scores es [B, T, T], un número por cada par de posiciones, y ahí está el cuadrado. La softmax convierte cada fila en pesos que suman uno, con los que la última línea calcula la media. Estas cuatro líneas dan lo mismo que F.scaled_dot_product_attention con una tolerancia de 1e-5.
Que sea una media tiene dos consecuencias. La primera, que cada coordenada de la salida queda entre el menor y el mayor de los values en esa coordenada, así que la atención nunca se sale de ese rango. La segunda, que si cambias una fila de la softmax por una fila one-hot (un solo 1 y ceros en el resto), la salida es exactamente el value de esa posición. El diccionario difuso contiene el exacto como caso límite.
Por qué la raíz cuadrada
La división por sqrt(d) parece un detalle, pero de ella depende que el modelo aprenda desde el primer paso o arranque saturado. El motivo es una cuenta que puedes comprobar en una consola.
El producto escalar de dos vectores aleatorios de anchura d con componentes de varianza 1 tiene varianza d. Lo he medido con d = 4096 y sale 4104. Si metes puntuaciones tan dispersas en una softmax, se satura: un peso se va casi a 1 y los demás, casi a 0. Con 1024 keys, el peso más grande de cada fila ha salido de media 0,976 y la fila se ha quedado en 0,09 bits de entropía de los 10 posibles. Eso es una búsqueda exacta disfrazada de difusa, y en plena inicialización. Un modelo entrenado también acaba afilando las filas, pero lo hace a propósito y desde un punto de partida del que puede aprender. Al dividir por sqrt(d), la varianza queda en 1,0 y la fila sube a 9,27 bits de entropía, con el peso más grande en 0,017. Con la pérdida calculada a partir de la media ponderada, el gradiente medio que llega a las puntuaciones es alrededor de un orden de magnitud mayor, entre 12 y 18 veces según la semilla. Si la pérdida se calcula a partir de un solo peso, el factor oscila entre unas pocas veces y muchos órdenes de magnitud, según las filas en las que le toque ganar a esa key. Una softmax saturada apenas tiene de qué aprender, porque su gradiente queda por debajo de 1e-6 en el 99,8 % de las puntuaciones y lo poco que hay se amontona en los contados casi empates.
Llama lo escribe una sola vez, self.scaling = self.head_dim**-0.5, en modeling_llama.py:226. Ahí la anchura que cuenta es la de la cabeza, sin que intervenga la del modelo.
Las cabezas son una vista y la máscara, un triángulo
La atención multicabeza consiste en ejecutar las cuatro líneas de arriba una vez por cabeza, H veces, sobre trozos de d. Trocear no cuesta nada: q.view(B, T, H, d // H).transpose(1, 2) da [B, H, T, d/H] y comparte memoria con q, porque el eje de las cabezas es solo un stride. Las puntuaciones pasan a ser [B, H, T, T], un cuadrado por cabeza. La cabeza 0 de la llamada en batch da lo mismo que la atención calculada solo sobre la cabeza 0. Lo que se gana con las cabezas son H juegos de pesos distintos sobre las mismas posiciones, cada uno en un trozo de anchura d/H (128 en Llama 3.1 8B, con 32 cabezas).
Es la máscara la que hace que un decodificador deje de promediar sobre todas las posiciones. Para predecir el token que sigue a la posición t, el modelo puede leer las posiciones 0..t y nada posterior, así que las puntuaciones por encima de la diagonal se ponen a -inf antes de la softmax y sus pesos quedan a cero. Lo he comprobado de la forma más directa, calculando la atención sobre ocho posiciones y reescribiendo luego las keys y los values de las posiciones 5 a 7. Las posiciones 0 a 4 han devuelto la misma salida con una tolerancia de 1e-6, mientras que sin la máscara han cambiado las ocho.
T = 12. Cada fila es una query y cada columna, una key; una celda encendida es una puntuación que llegará a la softmax. Elige una máscara y mira cuántas celdas quedan vivas.La query t lee las keys 0..t. Las puntuaciones por encima de la diagonal valen −inf antes de la softmax, así que sus pesos son cero.
Prueba las tres máscaras y mira el recuento. El tercer botón, el de la ventana, es la última palanca de la sección del coste. El triángulo se combina por suma con la máscara de padding de la parte 3. Si añades dos posiciones de padding a la izquierda de una secuencia, las dos primeras filas de query quedan enteras a -inf. En la softmax, una fila toda de -inf es 0/0, o sea, un NaN. Es la misma trampa de la parte anterior, ahora en la versión de cuatro líneas. La biblioteca transformers, en cambio, rellena su máscara eager con el mínimo del dtype (masking_utils.py:602-604) y esas filas salen uniformes. Así no dan error, pero siguen sin significar nada.
La ruta eager calcula las exponenciales y su suma en fp32 sea cual sea el dtype del modelo (dtype=torch.float32, en modeling_llama.py:208), y después vuelve a convertir la fila a ese dtype. Sumar el normalizador con los ocho bits significativos de bf16 sería ahorrar memoria donde no toca.
Lectura de LlamaAttention.forward
En transformers 5.17.0, modeling_llama.py:243-281 son esas cuatro líneas con las cabezas, la máscara y una caché, todo escrito por extenso. Abreviado y con las formas del modelo 8B, queda así:
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) # [B, 32, T, 128]key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) # [B, 8, T, 128]value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) # [B, 8, T, 128]query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)attn_output, attn_weights = attention_interface(self, query_states, key_states, value_states, attention_mask, scaling=self.scaling)attn_output = self.o_proj(attn_output.reshape(*input_shape, -1))Ahí hay tres cosas que no están en las cuatro líneas. La primera, past_key_values.update, es la caché. Al generar se guardan las keys y los values de todos los tokens anteriores, así que cada token nuevo calcula una sola fila de query y lee T keys guardadas.
La segunda es que la key y el value tienen 8 cabezas y la query, 32, de ahí que k_proj proyecte de 4096 a 1024. Es lo que se llama grouped-query attention (GQA), frente a una cabeza key-value por cada cabeza de query (MHA). Su razón de ser es la caché. La tercera, attention_interface, elige qué implementación se ejecuta, con eager_attention_forward como respaldo, en modeling_llama.py:191-213. Eager quiere decir que las cuatro líneas se ejecutan tal como están escritas, más repeat_kv, que en tiempo de ejecución copia las 8 cabezas de key hasta tener 32 para que cuadre la multiplicación de matrices. Si pones _attn_implementation a sdpa, la scaled_dot_product_attention de la primera sección, o a flash_attention_2, las mismas operaciones las calcula otro kernel (la rutina que hace la aritmética). La línea que aplica la rotación queda para la parte 7.
El coste lo ponen los datos y el código no lo delata
En software, el coste de una función se lee en su código: un bucle sobre la entrada es lineal, dos bucles anidados son cuadráticos, y lo sabes antes de ejecutarla. La atención es una sola función, su código son las cuatro líneas y en ellas no hay ningún bucle. El cuadrado está dentro de un único @ y nada en la forma del código delata que sea cuadrático. Su coste lo fija la T de la petición y crece con el cuadrado de las posiciones, por cabeza y por capa.
Con la forma de Llama 3.1 8B y en bf16, las puntuaciones de una secuencia en una capa ocupan 268 MB con 2048 tokens, 4,3 GB con 8192 y 1,1 TB con los 131 072 que declara la configuración. En el pico de memoria, la ruta eager mantiene tres cuadrados así (dos de ellos en fp32), que suman unas cinco veces esos bytes. Los pesos del modelo entero ocupan 16 GB. Si nadie llega a materializar esos 1,1 TB, es por el kernel.
[B, H, T, T], y la caché guarda las keys y los values de todos los tokens anteriores en todas las capas. Los pesos, 16 GB, aparecen dibujados como referencia de escala.Si llevas T a 16k con eager seleccionado, las puntuaciones de una sola capa ya superan los pesos del modelo entero, y con 131k son unas setenta veces más grandes. Al cambiar a sdpa, la cifra aparece tachada porque esa memoria nunca llega a reservarse.
He cronometrado las tres últimas líneas, con 8 cabezas de anchura 64, en la CPU de mi portátil. De 512 a 8192 tokens (dieciséis veces más entrada), la versión eager ha pasado de 1,9 ms a 284 ms, 150 veces más. En una capa lineal sobre la misma entrada, el tiempo ha subido de 0,23 ms a 2,5 ms, once veces más. Al doblar T de 4096 a 8192, el tiempo de la capa lineal se ha multiplicado por 1,95 y el de la atención por 4,3. Salvo por el ruido de la medición, es lo que predice el cuadrado.
La primera de las tres palancas que mueven esa cifra es el kernel. Con su backend flash o con el memory-efficient, F.scaled_dot_product_attention calcula los mismos números por bloques y nunca guarda entera la matriz T × T, cosa que su backend de respaldo, math, sí hace. Con 4096 tokens y en fp32, el profiler ha imputado 545 MB a las multiplicaciones de matrices de la versión eager (el cuadrado de 537 MB más su salida) y 8,5 MB a sdpa, que además ha ido 2,4 veces más rápido.
Flash attention lleva esa misma idea a la GPU, con una aritmética que sigue siendo cuadrática. Lo que cambia es lo que tiene que estar en memoria a la vez. La implementación se elige con un único campo de la configuración, _attn_implementation, que en 5.17.0 vale sdpa por defecto al cargar; eager solo lo tienes si lo pides. Si el batch lleva padding, sdpa sigue guardando la propia máscara, booleanos [B, 1, T, T] que ocupan 17 MB con 4096. Sin padding, transformers no pasa ninguna máscara y deja que el kernel aplique el triángulo (is_causal=True).
Las otras dos cambian lo que se calcula. La grouped-query attention reduce la caché. Con 8 cabezas key-value, el modelo 8B guarda 131 kB por token entre sus 32 capas, 17 GB para un contexto de 128k, ya algo más que los pesos. Con 32 cabezas serían 524 kB por token y 69 GB, cuatro veces los pesos. Una ventana deslizante, por su parte, limita hasta dónde puede leer hacia atrás una query. Llama 3.1 no tiene ninguna, Mistral 7B v0.1 tenía una, y Gemma 2 y 3 intercalan capas con ventana y capas globales. En transformers se escribe kv_idx > q_idx - sliding_window (masking_utils.py:100). Con 32 768 tokens, una ventana de 4096 deja 126 millones de celdas vivas frente a los 537 millones del triángulo causal. Mientras la secuencia no supera la ventana, el coste es un triángulo; cuando la supera, es una banda de anchura W que crece de forma lineal con T.
Para ser justos con el cuadrado, si en la pasada del prompt cuentas operaciones y dejas los bytes a un lado, las capas lineales del modelo 8B hacen más operaciones que la atención hasta unos 30 000 tokens, y hasta más o menos el doble si solo tienes en cuenta el triángulo que calcula un kernel causal. Con 8192 tokens, la atención supone alrededor de una quinta parte de las operaciones en coma flotante del bloque (sus FLOPs) si cuentas el cuadrado entero, y alrededor de una octava parte si cuentas solo el triángulo.
Como el cuadrado es un problema de memoria mucho antes de convertirse en uno de cómputo, la cuenta que hay que hacer antes de prometer una longitud de contexto es la de memoria.
Qué vigila el ingeniero de ML
Son cuatro cosas, la primera de ellas una multiplicación.
Calcula lo que ocupan [B, H, T, T] y la caché antes de prometer un contexto. El gráfico de arriba hace esa multiplicación, cabezas por T al cuadrado por dos bytes en cada capa. Súmale la caché por token multiplicada por la secuencia más larga y por el batch con el que vas a servir el modelo. El max_position_embeddings de la configuración es lo que el modelo admite tras estirar su rotación (aquí 131 072, una ventana de preentrenamiento de 8192 escalada y luego entrenada hasta 128k). Lo que cabe en tu GPU es otra cuenta.
Ten claro qué kernel se ejecuta. eager devuelve los pesos y materializa el cuadrado; sdpa y flash, no. Se elige en _attn_implementation y de ese campo depende que existan o no los 545 MB de antes. Una herramienta de depuración que necesite los pesos no puede sacarlos de sdpa, donde output_attentions=True devuelve una tupla vacía. Si para tenerlos cargas el modelo con attn_implementation="eager", el cuadrado se materializa en cada llamada.
Combina las máscaras a conciencia. La causal, la de padding y la de ventana se unen con un AND, que en la forma -inf es una suma. Una fila de query a la que ya no le queda nada que leer da NaN en esa forma y una fila uniforme en la de la biblioteca. Ninguno de los dos es un número que te sirva. La costumbre de la parte 3 vale ahora por partida doble: el índice del último token sale de la máscara y la máscara viaja con el batch.
Mide la latencia con el prompt más largo que vayas a aceptar. El tiempo de las capas lineales crece con T; el de la atención, con su cuadrado, y el de la lectura de la caché en cada token generado, con todo lo que esta ya contiene. Un p99 medido con 500 tokens no dice nada de lo que pasa con 8000. El salto de antes, de 1,9 ms a 284 ms, da una idea de la sorpresa que te espera.
Esto tampoco es trabajo de modelado. Una media ponderada cabe en cuatro líneas y sus pesos salen de tres matrices aprendidas. El cuadrado en el que viven es el primer coste de la serie que depende de la entrada sin que el código tenga nada que decir.