Cuánta memoria consume un modelo: cómo GQA, las ventanas y MLA cambian la caché KV

La situación resulta familiar: pones en marcha un modelo ya entrenado, los pesos caben sin problemas en la memoria de la GPU y el servidor falla por out-of-memory al llegar a la décima solicitud simultánea. La culpa no es del tamaño del modelo, sino de la caché KV. Su volumen depende de la arquitectura de la capa de atención, y eso ya no puedes elegirlo: la decisión se tomó durante el entrenamiento.

El checkpoint fija la arquitectura

La configuración y los pesos definen el número de cabezas de consulta y cabezas key-value, la dimensión de cada cabeza, el patrón de atención local y las dimensiones de la representación latente. Sustituir MHA por GQA después del entrenamiento supone modificar el modelo y, muy probablemente, habrá que recuperar calidad con entrenamiento adicional. Sustituir la atención completa por una ventana deslizante supone descartar contexto que el modelo está acostumbrado a ver.

Por eso, el procedimiento práctico consiste en abrir la configuración (en Hugging Face, con AutoConfig) y consultar el número de capas, num_attention_heads, num_key_value_heads, head_dim y los campos de sliding window. Una salvedad: el campo de la ventana puede figurar en la configuración y estar desactivado por un indicador adyacente, así que comprueba tanto el valor como el indicador. Algunos modelos alternan capas locales y globales, y el código de la implementación es siempre la referencia definitiva.

Cómo calcular la caché en una servilleta

La memoria persistente se mantiene entre los pasos de decodificación y determina cuántas solicitudes puedes atender a la vez. Los búferes temporales, como los bloques de puntuaciones de atención, pueden liberarse después de la capa. Para una caché estándar, la fórmula es: 2 × bytes por elemento × capas × tokens × número de cabezas KV × dimensión de cada cabeza. El factor 2 corresponde a las claves más los valores. Para un lote, multiplicamos por su tamaño.

Tomemos un modelo hipotético: 32 capas, 32 cabezas, head_dim 128, fp16 y un contexto de 8192 tokens. Con MHA clásico, una solicitud ocupa 4 GiB. Si el modelo usa GQA con 8 cabezas KV, ocupa 1 GiB. Con MQA y una sola cabeza KV, apenas 128 MiB. La diferencia entre «caben 4 solicitudes» y «caben 16» aparece sin cambiar el número de parámetros.

Las arquitecturas mixtas se calculan capa por capa. Supongamos que, en ese mismo modelo GQA de 32 capas, hay 24 locales con una ventana de 4096 y 8 de atención completa, con un contexto de 32 768 tokens. Usar atención completa en todas las capas daría 4 GiB; la arquitectura mixta, alrededor de 1,4 GiB. Pero la caché total sigue sin tener un límite fijo: cada token nuevo añade memoria en las capas de atención completa. Crece más despacio, pero no deja de crecer.

Qué variantes de atención hay

• MHA: cada cabeza de consulta tiene su propio par clave-valor. Es la caché más costosa y, durante la decodificación, también exige el mayor volumen de lectura de memoria.
• MQA y GQA: las cabezas de consulta comparten cabezas KV. La caché y los parámetros de las proyecciones K/V se reducen exactamente en el mismo factor que el número de cabezas KV.
• Ventana deslizante: la capa solo recuerda los últimos W tokens; la caché puede mantenerse en un búfer circular de tamaño fijo. Las posiciones se calculan a partir del índice lógico del token, no de la posición que ocupa en el búfer.
• MLA: en lugar de claves y valores completos, se almacena un vector latente comprimido. La caché es menor, pero la decodificación requiere más cálculos. Suele compensar porque el cuello de botella de la decodificación es el ancho de banda de memoria, no la capacidad de cálculo.

MLA no puede introducirse en la fórmula de MHA mediante num_key_value_heads: almacena otros tensores en la caché. Calcula a partir de lo que realmente guarda la implementación y utiliza kernels especializados.

Qué puede cambiar el runtime

Conviene separar cuatro cuestiones: qué tokens atienden a cuáles (atención completa o por ventanas), cuántas representaciones KV se han aprendido (MHA/GQA/MQA), qué se almacena en caché (tensores KV o una representación latente) y cómo se ejecuta. El checkpoint decide las tres primeras; el motor, la cuarta. FlashAttention calcula la atención exacta por bloques y evita escribir enormes matrices de puntuaciones en memoria: reduce los búferes temporales y el tráfico, mientras que la caché lógica y el número de cabezas KV permanecen iguales.

Paged attention almacena la caché en bloques que no tienen por qué ser contiguos y mejora la asignación y el uso compartido de memoria, pero no reduce por sí mismo el tamaño de la caché. El asignador debe respetar la política de expulsión del modelo: no puede descartar entradas antiguas de una capa de atención completa solo porque las capas vecinas usen ventanas.

Conclusión práctica

Antes de elegir un checkpoint, lee la configuración y estima la caché para el contexto y el tamaño de lote previstos. Si dos candidatos ofrecen una calidad similar, elige el que tenga cabezas KV o ventanas con menor coste de memoria: eso se traduce directamente en el número de solicitudes simultáneas por tarjeta. Y no olvides el padding: con lotes densos puedes estar pagando por posiciones vacías, algo que una implementación ragged evita.

El método tiene una limitación sencilla: el cálculo estima la memoria persistente, pero no incluye los pesos, las activaciones ni la sobrecarga del asignador. Las cifras anteriores son ilustrativas; para un modelo real, utiliza los valores de su configuración.