El KV cache tiene un problema, y no es la CPU
Cuando sirves un LLM con contextos largos, tarde o temprano te topas con la misma pared: la GPU se queda sin memoria mucho antes de quedarse sin cómputo. No es que el modelo no sepa hacer las cuentas; es que no tiene dónde guardarlas. Y el culpable casi siempre es el mismo: el KV cache, esa tabla de keys y values que crece con cada token para no recalcular la atención desde cero.
Lo que pocos miran es que ese cache no crece solo con la longitud del contexto. Crece también con el número de cabezas de atención. Ahí es donde entra la atención agrupada (grouped-query attention, GQA): un rediseño pequeño que recorta la memoria del cache varias veces sin que el modelo pierda apenas calidad. Es la razón de que los modelos grandes de los últimos años quepan donde antes no cabían.
De dónde sale el gasto: una copia por cabeza
La atención multi-cabeza clásica (multi-head attention, MHA) divide el trabajo en varias cabezas independientes. Cada cabeza aprende a fijarse en cosas distintas: una en la concordancia de género, otra en la referencia del pronombre, otra en la estructura de la frase. Para hacerlo, cada cabeza tiene su propia query, su propia key y su propio value.
El problema es el cache. Como las keys y los values de los tokens ya vistos no cambian, se guardan para reutilizarlos. Pero si tienes 32 cabezas, guardas 32 juegos de keys y values por cada token. Multiplica eso por 40.000 tokens de contexto y por decenas de peticiones simultáneas, y la cuenta de memoria se dispara.
La observación incómoda: buena parte de esa memoria es redundante. Muchas cabezas acaban atendiendo a patrones parecidos. ¿De verdad hace falta un juego completo de keys y values por cabeza?
Multi-query: el otro extremo
La primera respuesta fue radical. La atención multi-consulta (multi-query attention, MQA) mantiene una query por cabeza, pero hace que todas las cabezas compartan una sola key y un solo value. De 32 juegos en el cache pasas a uno. La memoria se desploma y el decode vuela, porque en cada paso lees muchísimos menos datos de la memoria de la GPU.
El precio es la calidad. Con una única representación de key/value para todas las cabezas, el modelo pierde parte de su capacidad de mirar el contexto desde ángulos distintos. En tareas exigentes se nota: la degradación es real y a veces el entrenamiento se vuelve inestable. MQA cambia demasiada precisión por memoria.
GQA: el punto medio que casi siempre gana
La atención agrupada es el término medio evidente en cuanto lo ves. En vez de un juego de key/value por cabeza (MHA) o uno para todas (MQA), agrupas las cabezas y cada grupo comparte un juego. Con 32 cabezas de query y 8 grupos, guardas 8 juegos de keys y values en lugar de 32. El cache se reduce a la cuarta parte y la calidad se mantiene a un pelo de la del MHA completo.
| Variante | Juegos de KV en cache | Memoria | Calidad |
|---|---|---|---|
| MHA (multi-cabeza) | uno por cabeza (p. ej. 32) | alta | referencia |
| MQA (multi-consulta) | uno para todas | mínima | cae en tareas duras |
| GQA (agrupada) | uno por grupo (p. ej. 8) | baja | ≈ MHA |
El motivo de que funcione es sencillo: agrupar conserva suficiente diversidad de representaciones para que el modelo siga viendo el contexto desde varios ángulos, mientras elimina la redundancia que MHA arrastra sin necesidad. Es un caso de libro de “el 80 % del beneficio con el 20 % del coste”.
Qué cambia en el código
En la práctica, el cambio es qué tensores compartes. Este pseudocódigo muestra la diferencia frente a MHA: las queries siguen siendo una por cabeza, pero varias cabezas apuntan al mismo grupo de key/value.
# n_cabezas_q: cabezas de query (p. ej. 32)
# n_grupos_kv: grupos de key/value (p. ej. 8) -> 4 cabezas por grupo
cabezas_por_grupo = n_cabezas_q // n_grupos_kv
def atencion_agrupada(x, cache_k, cache_v):
q = proyectar_query(x) # una query por cada una de las 32 cabezas
k, v = proyectar_kv(x) # solo 8 juegos: uno por grupo
cache_k.append(k) # el cache crece 4 veces más despacio
cache_v.append(v) # que en MHA
salida = []
for cabeza in range(n_cabezas_q):
grupo = cabeza // cabezas_por_grupo # 4 cabezas comparten el mismo grupo
# la query de esta cabeza atiende al key/value de SU grupo, no al suyo propio
salida.append(atender(q[cabeza], cache_k[grupo], cache_v[grupo]))
return combinar(salida)Fíjate en lo que no aparece: no hay 32 juegos de k y v. El bucle recorre las 32
queries, pero solo hay ocho representaciones de key/value que reutilizar. Esa
línea grupo = cabeza // cabezas_por_grupo es toda la idea.
Un detalle útil: MHA y MQA son casos límite de GQA. Si el número de grupos iguala al de cabezas, tienes MHA; si es uno, tienes MQA. GQA no es una tercera técnica exótica, sino el mando que regula cuánto compartes.
Cuándo te afecta a ti
No vas a implementar la atención a mano casi nunca, pero la decisión te alcanza igual.
Si sirves modelos y peleas con la memoria en contextos largos o con mucha concurrencia, mira cuántos grupos de key/value usa tu modelo. Un modelo con GQA te deja meter contextos más largos o más peticiones en paralelo con la misma GPU, porque el cache ocupa una fracción. Es dinero directo: más throughput por tarjeta.
Si eliges entre modelos de tamaño parecido, el número de grupos de key/value es una pista de coste de inferencia que rara vez sale en la ficha de marketing. Dos modelos de 70.000 millones de parámetros pueden costar muy distinto de servir según cómo esté configurada su atención.
Y si mides latencia, recuerda que el decode está limitado por el ancho de banda de memoria, no por el cómputo. Leer menos keys y values en cada paso es exactamente lo que hace que la generación vaya más rápida. GQA no es solo un ahorro de memoria: es también un acelerador de la fase que domina el tiempo de respuesta.
En una frase
La atención agrupada es lo que pasa cuando alguien mira el KV cache, se pregunta cuánto de esa memoria es redundante y encuentra el punto exacto en el que compartir deja de doler. La próxima vez que un modelo grande te quepa donde no esperabas o te responda más rápido de lo que su tamaño sugiere, hay bastantes probabilidades de que esté compartiendo cabezas por debajo. Merece la pena saber por qué.