Grouped-Query Attention (GQA) es una variante de la arquitectura Multi-Head Attention (MHA) utilizada en los modelos Transformer. Mientras que MHA asigna un conjunto independiente de 'key' y 'value heads' a cada 'query head', GQA agrupa un número 'G' de 'query heads' para compartir un único conjunto de 'key' y 'value heads'. Esto significa que, en lugar de tener 'N' 'key/value heads' para 'N' 'query heads' (como en MHA), GQA puede tener 'N' 'query heads' pero solo 'N/G' 'key/value heads'. Esta reducción en el número de 'key' y 'value heads' disminuye significativamente la cantidad de memoria necesaria para almacenar los tensores 'key' y 'value' en la caché durante la inferencia secuencial ('KV cache'), y reduce el ancho de banda de memoria requerido para acceder a ellos.
GQA ha sido adoptado por modelos de lenguaje grandes (LLMs) de vanguardia para mejorar la eficiencia de la inferencia. Un ejemplo prominente es el modelo Llama 2 de Meta, que utiliza GQA para optimizar el rendimiento. Otros modelos como Falcon 40B y 180B también han implementado GQA. Esta técnica es crucial en entornos de producción donde la latencia y el costo computacional son factores críticos. Las plataformas de inferencia de LLMs, como vLLM o TGI (Text Generation Inference), a menudo aprovechan GQA para maximizar el 'throughput' y minimizar el uso de memoria en GPUs, permitiendo servir a más usuarios simultáneamente o utilizar hardware menos costoso.
Para un Arquitecto de Sistemas, GQA es una optimización estratégica que impacta directamente la viabilidad económica y la escalabilidad de los servicios basados en LLMs. La reducción en el uso de memoria de la 'KV cache' significa que se pueden procesar secuencias de entrada más largas o lotes ('batches') más grandes en el mismo hardware, o que se puede utilizar hardware con menos memoria VRAM. Esto se traduce en menores costos operativos (CAPEX y OPEX) y una mayor capacidad de servicio. El 'trade-off' principal es una ligera disminución en la calidad del modelo en comparación con MHA puro en algunos escenarios, aunque la investigación ha demostrado que esta diferencia es a menudo insignificante para la mayoría de las aplicaciones prácticas. La decisión de adoptar GQA implica equilibrar la eficiencia computacional y el costo con los requisitos de rendimiento y calidad del modelo, siendo una elección casi obligatoria para la inferencia de LLMs a gran escala.