Grouped Query Attention (GQA)

Grouped Query Attention (GQA) é um design de atenção em que várias cabeças de consulta compartilham um conjunto menor de cabeças de chave e valor. Ele preserva as diversas perspectivas das consultas de múltiplas cabeças, ao mesmo tempo em que reduz a memória e o movimento de dados necessários para chaves e valores. O GQA é especialmente valioso durante a inferência autorregressiva, onde um modelo gera um token de cada vez e lê repetidamente os estados de atenção armazenados.

Como o GQA Funciona#

Dentro de um attention mechanism, cada token é projetado em três representações:

• Query: Descreve a informação que o token atual está procurando.
• Key: Descreve o que cada token disponível representa.
• Value: Contém a informação recuperada quando uma consulta corresponde a uma chave.

Na self-attention convencional, essas projeções são divididas em cabeças para que o modelo possa aprender diferentes relações em paralelo. O GQA mantém muitas cabeças de consulta, mas atribui grupos delas a cabeças de chave e valor compartilhadas. Por exemplo, uma camada de atenção pode usar oito cabeças de consulta e duas cabeças de chave-valor. Cada cabeça de chave-valor atende então a quatro cabeças de consulta.

O cálculo de atenção em si continua sendo a atenção de produto escalar em escala (scaled dot-product attention). O que muda é a disposição das cabeças e a quantidade de dados de chave-valor produzidos, armazenados e lidos. Frameworks como a PyTorch scaled dot-product attention, portanto, exigem que o número de cabeças de consulta seja divisível pelo número de cabeças de chave-valor quando o GQA está ativado. (docs.pytorch.org)

GQA vs. Atenção de Múltiplas Cabeças e de Múltiplas Consultas#

O GQA fica entre a atenção de múltiplas cabeças e a atenção de múltiplas consultas:

• Multi-head attention: Cada cabeça de consulta tem suas próprias cabeças de chave e valor. Isso oferece máxima independência de cabeças, mas cria o maior requisito de memória de chave-valor.
• Grouped query attention: Várias cabeças de consulta compartilham cada cabeça de chave-valor. Isso equilibra a capacidade de representação com a eficiência de memória.
• Multi-query attention: Todas as cabeças de consulta compartilham uma cabeça de chave e uma cabeça de valor. Isso minimiza o armazenamento de chave-valor, mas fornece menos diversidade de chave-valor.

Se uma camada tem 32 cabeças de consulta, a atenção de múltiplas cabeças também pode usar 32 cabeças de chave-valor, o GQA pode usar 8 e a atenção de múltiplas consultas usa 1. O GQA não é um substituto para a Transformer architecture mais ampla; ele é uma configuração possível dentro de uma camada de atenção do Transformer. A multi-head, multi-query, and grouped-query attention documentation da NVIDIA descreve essas variantes como diferindo principalmente em quantas cabeças de chave-valor atendem às cabeças de consulta. (nvidia.github.io)

Por que o GQA é Importante#

Durante a geração autorregressiva, um modelo armazena chaves e valores anteriores em um KV cache. Com a atenção de múltiplas cabeças padrão, cada cabeça de atenção contribui com chaves e valores armazenados em cache separados. Sequências longas, lotes grandes e muitas camadas de modelo podem, portanto, consumir memória substancial da GPU.

Como o GQA usa menos cabeças de chave-valor, seu cache é menor proporcionalmente à redução dessas cabeças. Isso pode melhorar:

• Capacidade de memória: Prompts mais longos ou mais solicitações simultâneas podem caber na memória disponível.
• Taxa de transferência de geração: Menos dados de chave-valor precisam ser lidos para cada token gerado.
• Inference latency: O tráfego de memória reduzido pode encurtar o tempo de processamento de token para token.
• Escalabilidade do Context window: Atender sequências longas torna-se mais prático.

Esses ganhos dependem do hardware, do comprimento da sequência, do tamanho do lote e do suporte do kernel. Os tempos de execução de produção ainda precisam de alocação de cache eficiente; o TensorRT-LLM KV cache system da NVIDIA, por exemplo, combina o suporte a GQA com cache baseado em blocos, reutilização e descarregamento. (nvidia.github.io)

Aplicações no Mundo Real#

Assistentes de contexto longo: Um assistente de codificação pode processar milhares de tokens de código-fonte antes de gerar uma resposta. O GQA reduz o estado de chave-valor em cache associado a esse histórico, permitindo que um servidor lide com arquivos mais longos ou mais usuários simultâneos sem aumentar proporcionalmente a memória da GPU.

Assistentes multimodais de imagem e vídeo: Um vision-language model pode representar patches de imagem ou quadros de vídeo como sequências longas de tokens. O GQA reduz a pressão do cache de atenção após a entrada desses tokens visuais no decodificador de linguagem, potencialmente deixando mais memória para imagens, quadros ou solicitações adicionais. Isso difere de um Vision Transformer padrão, onde a atenção pode processar todos os patches de imagem em paralelo sem cache KV autorregressivo.

Implementação Prática e Compromissos#

O PyTorch expõe o GQA através de enable_gqa=True. Este exemplo documentado do CUDA usa oito cabeças de consulta e duas cabeças de chave-valor compartilhadas:

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)

A saída retém oito cabeças de consulta, enquanto chaves e valores usam apenas duas cabeças. O PyTorch attention backend selector pode controlar qual kernel compatível realiza a operação, enquanto o PyTorch Transformer building-block guide fornece um contexto de implementação mais amplo. (docs.pytorch.org)

O GQA deve ser escolhido ao projetar ou adaptar um modelo; geralmente não é um interruptor em tempo de inferência para um ponto de verificação existente. Os desenvolvedores devem verificar o suporte do framework, a divisibilidade de cabeças, o comportamento numérico, o uso de memória e a qualidade da tarefa. Para implantação de visão computacional, fluxos de trabalho complementares, como a Ultralytics TensorRT integration e o Ultralytics benchmark mode, ajudam a medir se as otimizações de arquitetura e tempo de execução trazem melhorias significativas no hardware de destino.