MHA vs. MQA vs. GQA: Mecânica de Atenção e KV-Cache em LLMs

O Grouped-Query Attention (GQA) é o padrão de atenção dos modelos modernos de linguagem. Ele resolve o maior gargalo da inteligência artificial: o consumo de memória do KV-cache.

Modelos clássicos usam Multi-Head Attention (MHA), alocando pares exclusivos de chaves e valores para cada consulta. Isso consome dezenas de gigabytes de VRAM em contextos longos.

O Multi-Query Attention (MQA) comprime todas as consultas em um único par de chaves e valores. Essa abordagem economiza memória, mas degrada a capacidade de raciocínio.

O GQA atua como o meio-termo ideal da engenharia. Ele agrupa cabeças de consulta para compartilhar chaves e valores, cortando a memória em até 8x com acurácia quase intacta.

O que é atenção em LLMs e por que a geração satura a memória?

A geração de texto autorregressiva em modelos de linguagem opera no regime limitado por largura de banda de memória. A cada novo token gerado, a GPU transfere bilhões de pesos e todo o histórico do KV-cache para seus núcleos. Essa movimentação contínua esgota a memória antes do processador.

No processamento inicial do prompt (prefill), o modelo processa todos os tokens de entrada juntos. Ele executa multiplicações de matrizes densas (GEMM). A intensidade aritmética supera facilmente 150 FLOPs por byte transferido.

Na etapa de decodificação (decode), o modelo gera um único token por iteração. Ele executa operações matriz-vetor (GEMV). A intensidade aritmética despenca para a faixa de 1 a 8 FLOPs por byte.

O modelo Roofline formaliza esse teto de desempenho do hardware. Ele estabelece o ponto de equilíbrio (ridge point) de um acelerador moderno:

I_ridge = Capacidade_Computacional (FLOP/s) ÷ Largura_de_Banda_HBM (Bytes/s)

Uma GPU NVIDIA H100 SXM5 possui 3,35 TB/s de largura de banda e 989 TFLOPs em FP16. Seu ponto de equilíbrio atinge 295 FLOPs por byte.

Durante o decode, a intensidade cai para menos de 10 FLOPs por byte. Mais de 90% dos circuitos de cálculo ficam ociosos esperando dados da memória HBM.

O mecanismo de atenção seminal de Vaswani et al. (2017) salva vetores passados de chaves e valores. Isso evita recalcular projeções já computadas.

Esse buffer permanente na VRAM é o KV-cache. Quanto maior a janela de contexto, maior o volume de dados transferido pela GPU a cada fração de segundo.

O que é Multi-Head Attention (MHA) e por que ele se tornou insustentável?

No Multi-Head Attention tradicional, cada cabeça de consulta possui uma cabeça de chave e uma cabeça de valor exclusivas. Essa proporção direta de 1:1 garante que cada cabeça analise relações sintáticas e semânticas em projeções distintas. O modelo mantém máxima riqueza expressiva para raciocínio complexo.

O MHA projeta a ativação de entrada X em matrizes de consulta Q, chave K e valor V usando projeções lineares independentes:

Q = X × W_Q, K = X × W_K, V = X × W_V

Com H_q cabeças de consulta, o modelo aloca exatamente H_k = H_q e H_v = H_q. Cada cabeça opera com dimensão d_head.

Em sequências curtas de 512 tokens, essa arquitetura é eficiente. Contudo, em janelas de 32k a 128k tokens, o volume de tensores de chaves e valores armazenados em FP16 cresce linearmente com o contexto.

Considere um modelo com 32 camadas e 32 cabeças de atenção, com dimensão 128 por cabeça. Em precisão FP16 (2 bytes por valor), cada token consome:

S_token = 2 × 2 × 32 × 128 × 32 = 524.288 bytes ≈ 512 KB por token

Uma única requisição de 32.000 tokens consome 16,38 GB de memória VRAM apenas para armazenar o histórico de atenção. Um servidor com GPU de 80 GB não consegue atender simultaneamente mais de 4 usuários nesse contexto sem esgotar a memória física.

O que é Multi-Query Attention (MQA) e por que ele reduz a expressividade?

O Multi-Query Attention colapsa todas as cabeças de chave e valor em um único par compartilhado por todas as consultas. Proposto por Noam Shazeer (2019), o MQA reduz o tamanho do KV-cache pelo fator exato do número de cabeças.

Enquanto o modelo mantém H_q cabeças de consulta independentes, ele utiliza apenas H_k = 1 e H_v = 1 por camada. Todas as cabeças de consulta calculam atenção contra a mesmíssima sequência de chaves e valores.

Na decodificação, a GPU carrega apenas uma fração mínima de tensores da memória HBM. O throughput de geração de tokens aumenta em até 5x em servidores de produção.

A contrapartida é uma perda mensurável de capacidade expressiva. Como todas as consultas compartilham a mesma projeção, o modelo perde a habilidade de rastrear dependências sintáticas e entidades distintas em paralelo.

Em benchmarks de raciocínio lógico e código, modelos MQA sofrem degradação de qualidade. A perda é nítida quando comparada a modelos MHA de mesmo porte.

A engenharia de inteligência artificial demandava uma alternativa intermediária balanceada para conciliar velocidade e inteligência.

O que é Grouped-Query Attention (GQA) e como funciona a mecânica de grupos?

O Grouped-Query Attention divide as cabeças de consulta em grupos fixos que compartilham pares dedicados de chaves e valores. Proposto por Ainslie et al. (2023), o GQA une a qualidade representacional do MHA à eficiência computacional do MQA.

No GQA, as H_q cabeças de consulta são particionadas em G grupos uniformes. Cada grupo contém g = H_q ÷ G cabeças de consulta que compartilham um único par de chave e valor.

Se um modelo define H_q = 32 e G = 8, cada cabeça de chave atende a exatamente 4 cabeças de consulta. A razão de agrupamento é de 4:1.

Durante a operação de atenção, os tensores de chave e valor são replicados (broadcasted) para emparelhar com as consultas de seu respectivo grupo. Essa replicação ocorre nos registradores rápidos da GPU durante a computação, sem duplicar dados na memória HBM.

O Llama 3 (Meta AI, 2024) adotou GQA nativo em todas as versões. O Llama 3 8B adota 32 cabeças de consulta e 8 de chave e valor. A razão é 4:1.

As versões de 70B e 405B utilizam 64 cabeças de consulta e 8 de chave e valor. A razão atinge 8:1.

A mesma arquitetura está no Mistral 7B e no Gemma 2. Ela equipa também o Qwen 2.5. O GQA consolidou-se como padrão universal.

Comparativo técnico direto: MHA vs. MQA vs. GQA

A escolha da arquitetura de atenção dita a viabilidade econômica e operacional de um cluster de inferência de LLMs. Cada abordagem define um equilíbrio específico entre consumo de memória HBM, latência por token e fidelidade semântica.

A tabela a seguir consolida as características físicas, matemáticas e operacionais das três arquiteturas:

Modelos que adotam GQA mantêm pontuações equivalentes ao MHA em testes como MMLU, GSM8K e HumanEval. O ganho de throughput nos servidores de produção varia de 2,5x a 6x em cenários com lotes cheios.

Como calcular a memória VRAM do KV-Cache na prática?

O dimensionamento da memória de KV-cache exige calcular a quantidade total de parâmetros armazenados por token em todas as camadas do modelo. A fórmula exata em precisão FP16 ou BF16 (2 bytes por parâmetro) é:

S_KV = 2 × 2 × H_kv × d_head × L × seq_len × batch_size bytes

Nessa equação, o primeiro 2 representa o armazenamento conjunto de chaves e valores. O segundo 2 representa os 2 bytes de cada número em precisão de 16 bits.

O termo H_kv define o número de cabeças de chave e valor. O termo d_head é a dimensão de cada cabeça. O termo L define o total de camadas. Já seq_len é o contexto e batch_size é o lote.

Quando o runtime utiliza quantização de KV-cache em FP8 (1 byte por parâmetro), o multiplicador de precisão é reduzido de 2 para 1. Isso corta o consumo pela metade.

A tabela a seguir compara o consumo de VRAM entre MHA e GQA. Os valores cobrem o Llama 3 8B e o Llama 3 70B em diferentes contextos:

A redução de 8x no Llama 3 70B viabiliza manter o KV-cache em 4 GPUs A100 de 80 GB. Em MHA puro, seriam necessárias mais de 16 GPUs apenas para acomodar a memória do contexto.

Implementação em PyTorch: broadcasting de tensores no GQA

A execução do GQA em código exige expandir os tensores de chaves e valores para coincidir com a dimensão das consultas durante o produto escalar. Essa operação utiliza repetição por expansão de visualização (view broadcasting), evitando cópias físicas desnecessárias na memória principal.

O bloco a seguir ilustra a lógica canônica utilizada em bibliotecas como PyTorch e Transformers:

import torch
import torch.nn as nn
import math

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """Expande as cabeças de chave/valor para corresponder às cabeças de query."""
    if n_rep == 1:
        return x
    batch, num_kv_heads, seq_len, head_dim = x.shape
    return (
        x[:, :, None, :, :]
        .expand(batch, num_kv_heads, n_rep, seq_len, head_dim)
        .reshape(batch, num_kv_heads * n_rep, seq_len, head_dim)
    )

def grouped_query_attention(
    q: torch.Tensor,       # [B, H_q, S_q, D]
    k: torch.Tensor,       # [B, H_kv, S_kv, D]
    v: torch.Tensor,       # [B, H_kv, S_kv, D]
    mask: torch.Tensor = None
) -> torch.Tensor:
    B, H_q, S_q, D = q.shape
    H_kv = k.shape[1]
    n_rep = H_q // H_kv

    # Repetição eficiente nos registradores
    k_expanded = repeat_kv(k, n_rep)  # [B, H_q, S_kv, D]
    v_expanded = repeat_kv(v, n_rep)  # [B, H_q, S_kv, D]

    # Scaled Dot-Product Attention tradicional
    scores = torch.matmul(q, k_expanded.transpose(-2, -1)) / math.sqrt(D)
    if mask is not None:
        scores = scores + mask
    weights = torch.softmax(scores, dim=-1)
    return torch.matmul(weights, v_expanded)

Essa implementação calcula e armazena projeções de chave e valor apenas H_kv vezes. Isso economiza espaço crítico na memória VRAM.

Como servir modelos GQA com vLLM, PagedAttention e FlashDecoding?

A entrega eficiente de modelos baseados em GQA em servidores de produção exige runtimes que integrem gerenciamento de memória dinâmico e kernels otimizados. Motores de inferência modernos como vLLM e SGLang extraem o potencial máximo da arquitetura.

A tecnologia PagedAttention particiona o KV-cache em blocos virtuais de tamanho fixo. O conceito replica a memória virtual de sistemas operacionais.

Como o GQA usa menos cabeças por token, cada bloco paginado ocupa menos memória. O pool fixo de VRAM comporta muito mais blocos simultâneos.

Na decodificação, kernels como FlashAttention-2 e FlashDecoding particionam o histórico acumulado de tokens entre múltiplos multiprocessadores da GPU. A menor largura de banda exigida pelo GQA permite que todos os núcleos operem com saturação máxima sem esperar pela memória HBM.

Usar quantização FP8 no KV-cache com GQA reduz custos de nuvem em mais de 70%. A economia viabiliza projetos de grande porte.

Empresas migram cargas de clusters de 8 GPUs para servidores de 2 GPUs. Os mesmos acordos de nível de serviço são mantidos.

Perguntas Frequentes sobre MHA, MQA e GQA

O que significa a sigla GQA em inteligência artificial?

GQA significa Grouped-Query Attention, ou Atenção por Consulta Agrupada. Trata-se de uma arquitetura de rede neural que divide as cabeças de consulta em grupos fixos que compartilham pares de chaves e valores. O método equilibra a fidelidade semântica do MHA e a velocidade de execução do MQA.

Por que o Llama 3 adotou GQA em vez de MHA tradicional?

O Llama 3 adotou GQA para reduzir o footprint de memória do KV-cache e viabilizar contextos de 128.000 tokens em produção. Com GQA em razão 4:1 no modelo 8B e 8:1 no modelo 70B, a Meta reduziu a sobrecarga de memória sem comprometer o raciocínio. A mudança permite processar lotes muito maiores em servidores de inferência.

Qual a diferença entre MQA e GQA na arquitetura Transformer?

O Multi-Query Attention colapsa todas as consultas em um único par de chave e valor para toda a camada. O Grouped-Query Attention cria múltiplos grupos intermediários de consulta, onde cada grupo retém seu próprio par de chave e valor. O GQA evita a perda de expressividade do MQA preservando ganhos expressivos de memória.

Como o GQA reduz custos de infraestrutura em nuvem?

O GQA diminui o tamanho físico do KV-cache por usuário em até 8 vezes. Com menor consumo de memória por requisição, um mesmo servidor atende muito mais usuários concorrentes sem esgotar a VRAM. Isso permite utilizar instâncias de GPU menores e reduz a necessidade de alocar servidores adicionais.

É possível converter um modelo MHA para GQA sem retreinar do zero?

Sim, é possível converter um modelo existente através da técnica de uptraining proposta por Ainslie et al. O processo aplica agrupamento médio (mean-pooling) sobre as matrizes de projeção de chave e valor originais. Em seguida, o modelo passa por um ajuste fino rápido com cerca de 5% dos dados originais de pré-treino para recuperar a convergência.