Grouped Query Attention (GQA)

La atención de consultas agrupadas (GQA) es un diseño de atención en el que varios cabezales de consulta comparten un conjunto más pequeño de cabezales de clave y valor. Mantiene las distintas perspectivas de las consultas multi-cabezal, al tiempo que reduce la memoria y el movimiento de datos necesarios para las claves y los valores. GQA es especialmente útil durante la inferencia autorregresiva, en la que un modelo genera un token cada vez y lee repetidamente los estados de atención almacenados.

Cómo funciona GQA#

Dentro de un mecanismo de atención, cada token se proyecta en tres representaciones:

• Consulta: Describe la información que busca el token actual.
• Clave: Describe lo que representa cada token disponible.
• Valor: Contiene la información recuperada cuando una consulta coincide con una clave.

En la autoatención convencional, estas proyecciones se dividen en cabezales para que el modelo pueda aprender distintas relaciones en paralelo. GQA mantiene muchos cabezales de consulta, pero asigna grupos de ellos a cabezales de clave y valor compartidos. Por ejemplo, una capa de atención puede utilizar ocho cabezales de consulta y dos cabezales de clave y valor. Cada cabezal de clave y valor sirve entonces a cuatro cabezales de consulta.

El cálculo de atención en sí sigue siendo atención de producto punto escalado. Lo que cambia es la disposición de los cabezales y la cantidad de datos de clave y valor que se producen, almacenan y leen. Por tanto, los frameworks como la atención de producto punto escalado de PyTorch requieren que el número de cabezales de consulta sea divisible por el número de cabezales de clave y valor cuando GQA está activado. (docs.pytorch.org)

GQA frente a la atención multi-cabezal y multi-consulta#

GQA se sitúa entre la atención multi-cabezal y la atención multi-consulta:

• Atención multi-cabezal: Cada cabezal de consulta tiene sus propios cabezales de clave y valor. Esto ofrece la máxima independencia entre cabezales, pero crea el mayor requisito de memoria para claves y valores.
• Atención de consultas agrupadas: Varios cabezales de consulta comparten cada cabezal de clave y valor. Equilibra la capacidad de representación con la eficiencia de memoria.
• Atención multi-consulta: Todos los cabezales de consulta comparten un cabezal de clave y un cabezal de valor. Esto minimiza el almacenamiento de claves y valores, pero ofrece menos diversidad de claves y valores.

Si una capa tiene 32 cabezales de consulta, la atención multi-cabezal también puede utilizar 32 cabezales de clave y valor, GQA puede utilizar 8 y la atención multi-consulta utiliza 1. GQA no sustituye a la arquitectura Transformer general; es una posible configuración dentro de una capa de atención Transformer. La documentación de NVIDIA sobre la atención multi-cabezal, multi-consulta y de consultas agrupadas describe estas variantes como diferentes principalmente por el número de cabezales de clave y valor que sirven a los cabezales de consulta. (nvidia.github.io)

Por qué es importante GQA#

Durante la generación autorregresiva, un modelo almacena las claves y los valores anteriores en una caché KV. Con la atención multi-cabezal estándar, cada cabezal de atención aporta claves y valores almacenados en caché por separado. Por tanto, las secuencias largas, los lotes grandes y las numerosas capas del modelo pueden consumir una cantidad considerable de memoria de la GPU.

Como GQA utiliza menos cabezales de clave y valor, su caché es más pequeña en proporción a la reducción de esos cabezales. Esto puede mejorar:

• Capacidad de memoria: En la memoria disponible caben prompts más largos o más solicitudes simultáneas.
• Rendimiento de generación: Hay que leer menos datos de clave y valor por cada token generado.
• Latencia de inferencia: La reducción del tráfico de memoria puede acortar el tiempo de procesamiento entre tokens.
• Escalabilidad de la ventana de contexto: Resulta más práctico servir secuencias largas.

Estas mejoras dependen del hardware, la longitud de la secuencia, el tamaño del lote y la compatibilidad del kernel. Los runtimes de producción siguen necesitando una asignación eficiente de la caché; por ejemplo, el sistema de caché KV de TensorRT-LLM de NVIDIA combina la compatibilidad con GQA con el almacenamiento en caché basado en bloques, la reutilización y la descarga. (nvidia.github.io)

Aplicaciones en el mundo real#

Asistentes con contexto largo: Un asistente de programación puede procesar miles de tokens de código fuente antes de generar una respuesta. GQA reduce el estado de clave y valor almacenado en caché asociado a ese historial, lo que permite que un servidor gestione archivos más largos o más usuarios simultáneos sin aumentar proporcionalmente la memoria de la GPU.

Asistentes multimodales de imagen y vídeo: Un modelo de visión y lenguaje puede representar parches de imagen o fotogramas de vídeo como secuencias largas de tokens. GQA reduce la presión sobre la caché de atención después de que estos tokens visuales entren en el decodificador de lenguaje, lo que puede dejar más memoria para imágenes, fotogramas o solicitudes adicionales. Esto difiere de un Transformer de visión estándar, en el que la atención puede procesar todos los parches de imagen en paralelo sin almacenamiento en caché KV autorregresivo.

Implementación práctica y compromisos#

PyTorch expone GQA mediante enable_gqa=True. Este ejemplo documentado de CUDA utiliza ocho cabezales de consulta y dos cabezales de clave y valor compartidos:

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

assert torch.cuda.is_available(), "A CUDA device is required."

query = torch.randn(1, 8, 16, 32, device="cuda")
key = torch.randn(1, 2, 16, 32, device="cuda")
value = torch.randn(1, 2, 16, 32, device="cuda")

# Four query heads share each key-value head.
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        enable_gqa=True,
    )

print(output.shape)

La salida conserva ocho cabezales de consulta, mientras que las claves y los valores utilizan solo dos cabezales. El selector de backend de atención de PyTorch puede controlar qué kernel compatible ejecuta la operación, mientras que la guía de bloques de construcción de Transformer de PyTorch proporciona un contexto de implementación más amplio. (docs.pytorch.org)

GQA debe elegirse al diseñar o adaptar un modelo; por lo general, no es un interruptor de inferencia para un checkpoint existente. Los desarrolladores deben comprobar la compatibilidad del framework, la divisibilidad de los cabezales, el comportamiento numérico, el uso de memoria y la calidad de la tarea. Para el despliegue de visión artificial, flujos de trabajo complementarios como la integración de TensorRT de Ultralytics y el modo de benchmark de Ultralytics ayudan a medir si las optimizaciones arquitectónicas y del runtime ofrecen mejoras significativas en el hardware de destino.