MHA frente a MQA y GQA: cómo la atención cambia el uso de memoria KV

Traducción automática

Este artículo se tradujo automáticamente a partir de la versión original en inglés.

MHA asigna a cada cabeza de consulta su propia cabeza de clave y su propia cabeza de valor. MQA comparte una única cabeza de clave/valor entre todas las cabezas de consulta. GQA comparte una cabeza de clave/valor dentro de cada grupo de cabezas de consulta. Usar menos cabezas KV reduce la caché de atención convencional y los datos que la decodificación lee de ella.

Al estimar el tamaño de la caché, comprueba num_key_value_heads. Usar solo el número de cabezas de consulta puede sobreestimar considerablemente su tamaño.

Comparación directa del uso compartido de cabezas

Supongamos ocho cabezas de consulta y la misma dimensión por cabeza, número de capas, longitud de secuencia y precisión de la caché:

AtenciónCabezas de consultaCabezas KVCaché respecto a MHA
Atención multicabeza, MHA88100%
Atención multiconsulta, MQA8112.5%
Atención con consultas agrupadas, GQA8225%

Una cabeza KV incluye una cabeza de clave y una cabeza de valor. Las proyecciones de consulta siguen siendo independientes. El artículo de Shazeer sobre MQA y el artículo de Ainslie et al. sobre GQA definen estos cambios.

En Llama 3 70B, 64 cabezas de consulta comparten ocho cabezas KV: esto supone ocho veces menos datos KV sin procesar que en un modelo idéntico en todo lo demás con 64 cabezas KV. Llama 3.1 405B tiene 128 cabezas de consulta y ocho cabezas KV, lo que reduce el volumen dieciséis veces con la misma comparación. La tabla de arquitectura de Meta proporciona estas cifras.

Lo que la reducción no garantiza

Una KV cache ocho veces más pequeña no implica ocho veces menos memoria total de GPU ni una generación ocho veces más rápida. Los pesos, las activaciones, la ejecución de kernels y la comunicación siguen presentes. Las implementaciones distribuidas también pueden replicar cabezas KV entre particiones.

El artículo sobre GQA encontró un compromiso útil entre calidad y velocidad en los modelos de la familia T5 que evaluó. Convertir un checkpoint MHA ya entrenado a GQA es un cambio de arquitectura que requiere adaptación y evaluación; no es un ajuste de asignación de memoria.

Comprueba la configuración de atención del checkpoint, usa el número de cabezas KV en la fórmula de la caché y mide el rendimiento del servidor con las longitudes de contexto y la concurrencia que necesitas. Incluye la calidad en las tareas al comparar checkpoints.

Guía de ingeniería: GQA y MQA ilustra los patrones de uso compartido de cabezas.