Calculadora de memoria de atención
Datos de entrada
| Longitud de secuencia | 4.096 |
|---|---|
| Cabezas de atención | 32 |
| Tamaño de lote | 1 |
| Precisión | FP16 / BF16 (2 bytes) |
Calculadora de memoria de atención
Estima la memoria de la matriz de puntuaciones de atención materializada en la atención estándar de un transformer, a partir de la longitud de secuencia, el número de cabezas, el tamaño de lote y los bytes por elemento —el término cuadrático que FlashAttention elimina.
Datos de entrada
Carga de trabajo
Resultados
Introduce un valor para ver los resultados.
La memoria de atención
La atención estándar de un transformer compara cada token con todos los demás, construyendo una matriz de puntuaciones antes del softmax. Esa matriz es cuadrada en la longitud de secuencia, así que su memoria crece con el cuadrado del contexto —el famoso coste cuadrático de la atención. Esta calculadora estima el tamaño de esa matriz materializada a partir de la longitud de secuencia, el número de cabezas de atención, el tamaño de lote y los bytes usados por puntuación. Es precisamente la memoria que FlashAttention evita almacenar.
El término cuadrático
Para una secuencia de tokens, la atención forma una matriz de puntuaciones : la fila , columna guarda cuánto atiende el token al token . Almacenarla requiere memoria proporcional a . Esto es distinto de los pesos del modelo o de la caché clave-valor, que crecen linealmente con el contexto —la matriz de puntuaciones crece con su cuadrado, así que con contextos largos se convierte en el buffer intermedio dominante y en la razón de que la atención ingenua se quede sin memoria.
La fórmula
Cada cabeza construye su propia matriz de puntuaciones y cada secuencia del lote arrastra su propia copia, así que la memoria en bytes es
Abytes=B⋅H⋅s2⋅edonde es el tamaño de lote, el número de cabezas de atención, la longitud de secuencia y los bytes por puntuación almacenada. Dividir entre da gigabytes. El cuadrado sobre es lo que encarece el contexto largo: el número de cabezas y el lote solo multiplican linealmente.
Qué es esta memoria
La cifra mide la matriz de puntuaciones materializada que la atención estándar escribe en memoria entre el producto consulta-clave y la suma ponderada por softmax sobre los valores. Es el buffer de activación que una implementación ingenua debe retener. Si persiste o no depende del contexto: en un paso hacia adelante simple es transitorio y se puede reutilizar entre capas, mientras que los frameworks de entrenamiento pueden conservar una copia por capa para el paso hacia atrás, salvo que intervengan el checkpointing o un núcleo fusionado.
FlashAttention
FlashAttention calcula la salida idéntica sin escribir nunca la matriz completa . Recorre las claves y los valores en bloques pequeños, manteniendo estadísticas de softmax en curso para necesitar solo un bloque en memoria cada vez, lo que convierte el almacenamiento de cuadrático a lineal en la longitud de secuencia. La matriz que dimensiona esta calculadora es exactamente lo que FlashAttention se niega a almacenar —así que el resultado de aquí es un buen indicador de la memoria que un núcleo de atención fusionado ahorra a una longitud de contexto dada.
Ejemplo resuelto
Tomemos una única secuencia de 4096 tokens con 32 cabezas de atención en precisión de 16 bits:
Abytes=1×32×40962×2=1073741824unos 1,07 GB para la matriz de puntuaciones de una capa. Al duplicar el contexto a 8192 tokens, el término cuadrático toma el control: la misma expresión da aproximadamente 4,29 GB, cuatro veces más para el doble de longitud. El lado clave-valor del mismo presupuesto de atención se trata en la Calculadora del tamaño de la caché KV, y la porción de los pesos del modelo en la Calculadora de VRAM para inferencia de LLM.
Preguntas frecuentes (FAQ)
¿Por qué la memoria de atención es cuadrática en la longitud de secuencia?
La atención compara cada token con todos los demás, produciendo una matriz de puntuaciones cuyas dos dimensiones son la longitud de secuencia. Almacenar esa matriz completa requiere, por tanto, memoria proporcional al cuadrado de la longitud de secuencia.
Duplicar el contexto cuadruplica la matriz de puntuaciones, por lo que la atención estándar se vuelve limitada por la memoria con contextos largos —el término cuadrático supera al coste lineal de los pesos y las activaciones.
¿Cómo reduce FlashAttention esta memoria?
FlashAttention calcula el mismo resultado sin materializar nunca la matriz de puntuaciones completa. Recorre la atención por bloques pequeños de claves y valores, manteniendo estadísticas de softmax en curso para retener solo un bloque cada vez.
La matriz cuadrática que mide esta calculadora es exactamente la memoria que FlashAttention evita almacenar, lo que le permite escalar a secuencias mucho más largas. La aritmética no cambia; solo el almacenamiento intermedio se reduce de cuadrático a lineal.
¿Esta memoria es por capa o para todo el modelo?
Depende de la implementación. La cifra de aquí es el tamaño de la matriz de puntuaciones de una sola capa. En un paso hacia adelante ingenuo ese buffer se puede liberar y reutilizar entre capas, así que es transitorio y no se multiplica por el número de capas.
Durante el entrenamiento, en cambio, los frameworks pueden retener los tensores de atención por capa para el paso hacia atrás, salvo que se use checkpointing de gradiente o FlashAttention, en cuyo caso el total puede escalar con el número de capas.
Aviso legal
Esta estimación cubre solo la matriz de puntuaciones materializada a la precisión indicada; excluye los tensores de consulta, clave y valor, la proyección de salida y otras activaciones. Describe la atención estándar —los núcleos como FlashAttention no almacenan esta matriz en absoluto.